[utils] Improve strip_or_none

meunierd · Feb 13, 2020 · b379ac7 · b379ac7
1 parent 648db9e
commit b379ac7
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/test/test_utils.py b/test/test_utils.py
@@ -73,6 +73,7 @@
     smuggle_url,
     str_to_int,
     strip_jsonp,
+    strip_or_none,
     timeconvert,
     unescapeHTML,
     unified_strdate,
@@ -752,6 +753,18 @@ def test_strip_jsonp(self):
         d = json.loads(stripped)
         self.assertEqual(d, {'status': 'success'})
 
+    def test_strip_or_none(self):
+        self.assertEqual(strip_or_none(' abc'), 'abc')
+        self.assertEqual(strip_or_none('abc '), 'abc')
+        self.assertEqual(strip_or_none(' abc '), 'abc')
+        self.assertEqual(strip_or_none('\tabc\t'), 'abc')
+        self.assertEqual(strip_or_none('\n\tabc\n\t'), 'abc')
+        self.assertEqual(strip_or_none('abc'), 'abc')
+        self.assertEqual(strip_or_none(''), '')
+        self.assertEqual(strip_or_none(None), None)
+        self.assertEqual(strip_or_none(42), None)
+        self.assertEqual(strip_or_none([]), None)
+
     def test_uppercase_escape(self):
         self.assertEqual(uppercase_escape('aä'), 'aä')
         self.assertEqual(uppercase_escape('\\U0001d550'), '𝕐')

diff --git a/youtube_dl/utils.py b/youtube_dl/utils.py
@@ -1951,8 +1951,8 @@ def bool_or_none(v, default=None):
     return v if isinstance(v, bool) else default
 
 
-def strip_or_none(v):
-    return None if v is None else v.strip()
+def strip_or_none(v, default=None):
+    return v.strip() if isinstance(v, compat_str) else default
 
 
 def url_or_none(url):