[utils] Allow None in remove_{start,end}
authorSergey M․ <dstftw@gmail.com>
Wed, 18 May 2016 22:31:30 +0000 (04:31 +0600)
committerSergey M․ <dstftw@gmail.com>
Wed, 18 May 2016 22:31:30 +0000 (04:31 +0600)
test/test_utils.py
youtube_dl/utils.py

index 520d32ff505e66bc577004440a4082bc643d5249..a697232a8b8da8b45ca54338338a1c67dbf8977d 100644 (file)
@@ -50,6 +50,8 @@ from youtube_dl.utils import (
     sanitize_path,
     prepend_extension,
     replace_extension,
+    remove_start,
+    remove_end,
     remove_quotes,
     shell_quote,
     smuggle_url,
@@ -215,6 +217,16 @@ class TestUtil(unittest.TestCase):
         self.assertEqual(replace_extension('.abc', 'temp'), '.abc.temp')
         self.assertEqual(replace_extension('.abc.ext', 'temp'), '.abc.temp')
 
+    def test_remove_start(self):
+        self.assertEqual(remove_start(None, 'A - '), None)
+        self.assertEqual(remove_start('A - B', 'A - '), 'B')
+        self.assertEqual(remove_start('B - A', 'A - '), 'B - A')
+
+    def test_remove_end(self):
+        self.assertEqual(remove_end(None, ' - B'), None)
+        self.assertEqual(remove_end('A - B', ' - B'), 'A')
+        self.assertEqual(remove_end('B - A', ' - B'), 'B - A')
+
     def test_remove_quotes(self):
         self.assertEqual(remove_quotes(None), None)
         self.assertEqual(remove_quotes('"'), '"')
index ac60ba18cd9fd9c5ecba155af3752458d0729766..5301d07404b4d83c4df5c99dd9a69f577b87f091 100644 (file)
@@ -1549,15 +1549,11 @@ def setproctitle(title):
 
 
 def remove_start(s, start):
-    if s.startswith(start):
-        return s[len(start):]
-    return s
+    return s[len(start):] if s is not None and s.startswith(start) else s
 
 
 def remove_end(s, end):
-    if s.endswith(end):
-        return s[:-len(end)]
-    return s
+    return s[:-len(end)] if s is not None and s.endswith(end) else s
 
 
 def remove_quotes(s):