From 0cdef1d74062239e8bae475b2f04274a92821f39 Mon Sep 17 00:00:00 2001
From: Mattijs Ugen <144798+akaIDIOT@users.noreply.github.com>
Date: Sun, 15 Dec 2019 17:59:22 +0100
Subject: [PATCH] Create timezone-aware datetimes when parsed as such (#163)

* On load, now use aware datetimes if possible

On loading data, if timestamps have an ISO "+HH:MM" UTC offset then the resultant datetime is converted to UTC.  This change adds that timezone information to the datetime objects.

Importantly, this addresses a Django warning (and potential error) that appears when using both YAML fixtures in a timezone-aware project.  It was raised as a Django issue (https://code.djangoproject.com/ticket/18867), but subsequently closed because the Django devs felt that this is a PyYAML problem.

* Create timezone-aware datetime in timezone from data

* Create timezone-aware datetime in timezone from data for python2

* Define better timezone implementation for python2

* Handle timezone "Z" for python 3

* Handle timezone "Z" for python 2

* Fix code structure for Python 3

Call datetime.datetime constructor once at return.

* Fix code structure for Python 2

Call datetime.datetime constructor once at return.
---
 lib/yaml/constructor.py  | 36 ++++++++++++++++++++++++++++++------
 lib3/yaml/constructor.py | 11 ++++++-----
 2 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/lib/yaml/constructor.py b/lib/yaml/constructor.py
index 08bdabfb..568613b7 100644
--- a/lib/yaml/constructor.py
+++ b/lib/yaml/constructor.py
@@ -18,6 +18,29 @@
 class ConstructorError(MarkedYAMLError):
     pass
 
+
+class timezone(datetime.tzinfo):
+    def __init__(self, offset):
+        self._offset = offset
+        seconds = abs(offset).total_seconds()
+        self._name = '%s%02d:%02d' % (
+            '-' if offset.days < 0 else '+',
+            seconds // 3600,
+            seconds % 3600 // 60
+        )
+
+    def tzname(self, dt=None):
+        return self._name
+
+    def utcoffset(self, dt=None):
+        return self._offset
+
+    def dst(self, dt=None):
+        return datetime.timedelta(0)
+
+    __repr__ = __str__ = tzname
+
+
 class BaseConstructor(object):
 
     yaml_constructors = {}
@@ -293,7 +316,7 @@ def construct_yaml_binary(self, node):
             return str(value).decode('base64')
         except (binascii.Error, UnicodeEncodeError), exc:
             raise ConstructorError(None, None,
-                    "failed to decode base64 data: %s" % exc, node.start_mark) 
+                    "failed to decode base64 data: %s" % exc, node.start_mark)
 
     timestamp_regexp = re.compile(
             ur'''^(?P<year>[0-9][0-9][0-9][0-9])
@@ -320,22 +343,23 @@ def construct_yaml_timestamp(self, node):
         minute = int(values['minute'])
         second = int(values['second'])
         fraction = 0
+        tzinfo = None
         if values['fraction']:
             fraction = values['fraction'][:6]
             while len(fraction) < 6:
                 fraction += '0'
             fraction = int(fraction)
-        delta = None
         if values['tz_sign']:
             tz_hour = int(values['tz_hour'])
             tz_minute = int(values['tz_minute'] or 0)
             delta = datetime.timedelta(hours=tz_hour, minutes=tz_minute)
             if values['tz_sign'] == '-':
                 delta = -delta
-        data = datetime.datetime(year, month, day, hour, minute, second, fraction)
-        if delta:
-            data -= delta
-        return data
+            tzinfo = timezone(delta)
+        elif values['tz']:
+            tzinfo = timezone(datetime.timedelta(0))
+        return datetime.datetime(year, month, day, hour, minute, second, fraction,
+                                 tzinfo=tzinfo)
 
     def construct_yaml_omap(self, node):
         # Note: we do not check for duplicate keys, because it's too
diff --git a/lib3/yaml/constructor.py b/lib3/yaml/constructor.py
index ac15aba8..cd9167ea 100644
--- a/lib3/yaml/constructor.py
+++ b/lib3/yaml/constructor.py
@@ -324,22 +324,23 @@ def construct_yaml_timestamp(self, node):
         minute = int(values['minute'])
         second = int(values['second'])
         fraction = 0
+        tzinfo = None
         if values['fraction']:
             fraction = values['fraction'][:6]
             while len(fraction) < 6:
                 fraction += '0'
             fraction = int(fraction)
-        delta = None
         if values['tz_sign']:
             tz_hour = int(values['tz_hour'])
             tz_minute = int(values['tz_minute'] or 0)
             delta = datetime.timedelta(hours=tz_hour, minutes=tz_minute)
             if values['tz_sign'] == '-':
                 delta = -delta
-        data = datetime.datetime(year, month, day, hour, minute, second, fraction)
-        if delta:
-            data -= delta
-        return data
+            tzinfo = datetime.timezone(delta)
+        elif values['tz']:
+            tzinfo = datetime.timezone.utc
+        return datetime.datetime(year, month, day, hour, minute, second, fraction,
+                                 tzinfo=tzinfo)
 
     def construct_yaml_omap(self, node):
         # Note: we do not check for duplicate keys, because it's too