Introduce validate_filename()

danny0838 · danny0838 · commit 978ace65cf22 · 2020-10-12T22:33:30.000+08:00
diff --git a/test/test_util.py b/test/test_util.py
@@ -111,6 +111,32 @@ def test_id_to_datetime_legacy(self):
             datetime(2020, 1, 2, 3, 4, 5))
         self.assertIsNone(util.id_to_datetime_legacy('20200102'), None)
 
+    def test_validate_filename(self):
+        self.assertEqual(
+            util.validate_filename(''),
+            '_')
+        self.assertEqual(
+            util.validate_filename('.wsb'),
+            '_.wsb')
+        self.assertEqual(
+            util.validate_filename('foo.'),
+            'foo')
+        self.assertEqual(
+            util.validate_filename('  wsb  '),
+            'wsb')
+        self.assertEqual(
+            util.validate_filename(''.join(chr(i) for i in range(0x80))),
+            "!_#$%&'()_+,-._0123456789_;(=)_@ABCDEFGHIJKLMNOPQRSTUVWXYZ[_]^_`abcdefghijklmnopqrstuvwxyz{_}-")
+        self.assertEqual(
+            util.validate_filename('\u0080中文𠀀'),
+            '\u0080中文𠀀')
+        self.assertEqual(
+            util.validate_filename(''.join(chr(i) for i in range(0x80)), force_ascii=True),
+            "!_#$%25&'()_+,-._0123456789_;(=)_@ABCDEFGHIJKLMNOPQRSTUVWXYZ[_]^_`abcdefghijklmnopqrstuvwxyz{_}-")
+        self.assertEqual(
+            util.validate_filename('\u0080中文𠀀', force_ascii=True),
+            '%C2%80%E4%B8%AD%E6%96%87%F0%A0%80%80')
+
     def test_crop(self):
         self.assertEqual(util.crop('dummy text', 10), 'dummy text')
         self.assertEqual(util.crop('dummy text', 9), 'dummy ...')
diff --git a/webscrapbook/util.py b/webscrapbook/util.py
@@ -17,7 +17,7 @@
 import binascii
 import codecs
 from base64 import b64decode
-from urllib.parse import unquote_to_bytes
+from urllib.parse import quote, unquote_to_bytes
 from urllib.request import pathname2url
 from ipaddress import IPv6Address, AddressValueError
 from datetime import datetime, timezone
@@ -159,6 +159,39 @@ def id_to_datetime_legacy(id):
     return None
 
 
+def validate_filename(filename, force_ascii=False):
+    """Transliterates the given string to be a safe filename
+
+    See also: scrapbook.validateFilename of WebScrapBook.
+    """
+    fn = filename
+
+    # control chars are bad for filename
+    fn = re.sub(r'[\x00-\x1F\x7F]+', '', fn)
+
+    # leading/trailing spaces and dots are not allowed on Windows
+    fn = re.sub(r'^\.', '_.', fn)
+    fn = re.sub(r'^ +', '', fn)
+    fn = re.sub(r'[. ]+$', '', fn)
+
+    # bad chars on most OS
+    fn = re.sub(r'[:"?*\\/|]', '_', fn)
+
+    # bad chars on Windows, replace with adequate direction
+    fn = fn.replace('<', '(').replace('>', ')')
+
+    # "~" is not allowed by browser.downloads
+    fn = fn.replace('~', '-')
+
+    if force_ascii:
+        fn = quote(fn, safe="""!_#$&'()*+,-./:;<=>?@[\\]^_`{|}~""")
+
+    # prevent empty filename
+    fn = fn or "_"
+
+    return fn
+
+
 #########################################################################
 # String handling
 #########################################################################
