utils: add parse_duration(), which takes in a duration string and returns the equiv. amt of seconds (#504)

test/test_utils.py

@@ -149,6 +149,44 @@ class UtilsTestCase(unittest.TestCase):
         self.assertEqual(utils.get_hostname_type("1:"), 0)
         self.assertEqual(utils.get_hostname_type(":5"), 0)
 
+    def test_parse_duration(self):
+        # Base case: simple number
+        self.assertEqual(utils.parse_duration("0"), 0)
+        self.assertEqual(utils.parse_duration("256"), 256)
+
+        # Not valid: not a positive integer
+        with self.assertRaises(ValueError):
+            utils.parse_duration("-5")
+            utils.parse_duration("3.1416")
+
+        # Not valid: wrong units or nonsense
+        with self.assertRaises(ValueError):
+            utils.parse_duration("")
+            utils.parse_duration("3j")
+            utils.parse_duration("5h6")  # stray number at end
+            utils.parse_duration("5h3k")
+            utils.parse_duration(" 6d ")
+            utils.parse_duration("6.6d")  # we don't support monster math
+            utils.parse_duration("zzzzzdstwataw")
+            utils.parse_duration("3asdfjkl;")
+
+        # Test all supported units
+        self.assertEqual(utils.parse_duration("3s"), 3)
+        self.assertEqual(utils.parse_duration("1m"), 60)
+        self.assertEqual(utils.parse_duration("9h"), 9 * 60 * 60)
+        self.assertEqual(utils.parse_duration("15d"), 15 * 24 * 60 * 60)
+        self.assertEqual(utils.parse_duration("3w"), 3 * 7 * 24 * 60 * 60)
+
+        # Composites
+        self.assertEqual(utils.parse_duration("6m10s"), 6 * 60 + 10)
+        self.assertEqual(utils.parse_duration("1d5h"), ((24+5) * 60 * 60))
+        self.assertEqual(utils.parse_duration("2d3m4s"), (48 * 60 * 60 + 3 * 60 + 4))
+
+        # Not valid: wrong order of units
+        with self.assertRaises(ValueError):
+            utils.parse_duration("4s3d")
+            utils.parse_duration("1m5w")
+
 
 if __name__ == '__main__':
     unittest.main()

utils.py

@@ -768,3 +768,48 @@ def get_hostname_type(address):
             return 2
         else:
             raise ValueError("Got unknown value %r from ipaddress.ip_address()" % address)
+
+_duration_re = re.compile(r"^((?P<week>\d+)w)?((?P<day>\d+)d)?((?P<hour>\d+)h)?((?P<minute>\d+)m)?((?P<second>\d+)s)?$")
+def parse_duration(text):
+    """
+    Takes in a duration string and returns the equivalent amount of seconds.
+
+    Time strings are in the following format:
+    - '123'         => 123 seconds
+                       (positive integers are treated as # of seconds)
+    - '1w2d3h4m5s'  => 1 week, 2 days, 3 hours, 4 minutes, and 5 seconds
+                       (must be in decreasing order by unit)
+    - '72h'         => 72 hours
+    - '1h5s'        => 1 hour and 5 seconds
+    and so on...
+    """
+    # If we get an already valid number, just return it
+    if text.isdigit():
+        return int(text)
+
+    match = _duration_re.match(text)
+    if not match:
+        raise ValueError("Failed to parse duration string %r" % text)
+    result = 0
+    matched = 0
+
+    if match.group('week'):
+        result += int(match.group('week'))   *  7 * 24 * 60 * 60
+        matched += 1
+    if match.group('day'):
+        result += int(match.group('day'))    * 24 * 60 * 60
+        matched += 1
+    if match.group('hour'):
+        result += int(match.group('hour'))   * 60 * 60
+        matched += 1
+    if match.group('minute'):
+        result += int(match.group('minute')) * 60
+        matched += 1
+    if match.group('second'):
+        result += int(match.group('second'))
+        matched += 1
+
+    if not matched:
+        raise ValueError("Failed to parse duration string %r" % text)
+
+    return result