datetimeex.construct_yaml_timestamp: preserve timezone

ioex/datetimeex.py

@@ -10,11 +10,20 @@ def construct_yaml_timestamp(loader, node):
     loaded_dt = loader.construct_yaml_timestamp(node)
     if type(loaded_dt) is datetime.datetime and loaded_dt.tzinfo is None:
         timezone_match = re.search(
-            ur'(Z|[\+-]\d{2}:\d{2})$',
+            ur'(Z|(?P<sign>[\+-])(?P<h>\d{2}):(?P<m>\d{2}))$',
             loader.construct_yaml_str(node),
             )
         if timezone_match:
+            # loader.construct_yaml_timestamp converts to UTC
             loaded_dt = loaded_dt.replace(tzinfo = dateutil.tz.tz.tzutc())
+            timezone_attr = timezone_match.groupdict()
+            if timezone_attr['h']:
+                timezone = dateutil.tz.tz.tzoffset(
+                        name = timezone_match.group(0),
+                        offset = (int(timezone_attr['h']) * 60 + int(timezone_attr['m'])) * 60
+                            * (-1 if timezone_attr['sign'] == '-' else 1),
+                        )
+                loaded_dt = loaded_dt.astimezone(timezone)
     return loaded_dt
 
 def add_yaml_timestamp_constructor(loader, tag = u'tag:yaml.org,2002:timestamp'):

tests/datetimeex/test_yaml_timestamp_constructor.py

@@ -16,11 +16,14 @@ import pytz
     ['2016-01-14 13:50:04+01:00', pytz.timezone('Europe/Vienna').localize(datetime.datetime(2016, 1, 14, 13, 50, 4, 0))],
     ['2016-07-14 13:50:04+02:00', pytz.timezone('Europe/Vienna').localize(datetime.datetime(2016, 7, 14, 13, 50, 4, 0))],
     ['2016-07-14 13:50:04+02:00', datetime.datetime(2016, 7, 14, 13, 50, 4, 0, tzinfo = dateutil.tz.tz.tzoffset('Vienna', 2 * 60 * 60))],
+    ['2016-07-14 13:50:04-07:00', pytz.timezone('US/Pacific').localize(datetime.datetime(2016, 7, 14, 13, 50, 4, 0))],
     ])
 def test_from_yaml(yaml_string, expected_timestamp, loader):
     loader_copy = copy.deepcopy(loader)
     ioex.datetimeex.add_yaml_timestamp_constructor(loader_copy)
-    assert yaml.load(yaml_string, Loader = loader_copy) == expected_timestamp
+    loaded_timestamp = yaml.load(yaml_string, Loader = loader_copy)
+    assert loaded_timestamp == expected_timestamp
+    assert loaded_timestamp.utcoffset() == expected_timestamp.utcoffset()
 
 @pytest.mark.parametrize(('yaml_string', 'tag', 'expected_timestamp'), [
     ['!without_timezone 2016-07-14 13:50:04', '!without_timezone', datetime.datetime(2016, 7, 14, 13, 50, 4, 0)],