gh-53502: add a new option aware_datetime in plistlib to loads or dumps aware datetime. (#113363)

* add options to loads and dumps aware datetime in plistlib

4 files changed, 134 insertions, 23 deletions

diff --git a/Doc/library/plistlib.rst b/Doc/library/plistlib.rst
index 732ef35..10f1a48 100644
--- a/Doc/library/plistlib.rst
+++ b/Doc/library/plistlib.rst
@@ -52,7 +52,7 @@ or :class:`datetime.datetime` objects.
 
 This module defines the following functions:
 
-.. function:: load(fp, *, fmt=None, dict_type=dict)
+.. function:: load(fp, *, fmt=None, dict_type=dict, aware_datetime=False)
 
    Read a plist file. *fp* should be a readable and binary file object.
    Return the unpacked root object (which usually is a
@@ -69,6 +69,10 @@ This module defines the following functions:
    The *dict_type* is the type used for dictionaries that are read from the
    plist file.
 
+   When *aware_datetime* is true, fields with type ``datetime.datetime`` will
+   be created as :ref:`aware object <datetime-naive-aware>`, with
+   :attr:`!tzinfo` as :attr:`datetime.UTC`.
+
    XML data for the :data:`FMT_XML` format is parsed using the Expat parser
    from :mod:`xml.parsers.expat` -- see its documentation for possible
    exceptions on ill-formed XML.  Unknown elements will simply be ignored
@@ -79,8 +83,11 @@ This module defines the following functions:
 
    .. versionadded:: 3.4
 
+   .. versionchanged:: 3.13
+      The keyword-only parameter *aware_datetime* has been added.
+
 
-.. function:: loads(data, *, fmt=None, dict_type=dict)
+.. function:: loads(data, *, fmt=None, dict_type=dict, aware_datetime=False)
 
    Load a plist from a bytes object. See :func:`load` for an explanation of
    the keyword arguments.
@@ -88,7 +95,7 @@ This module defines the following functions:
    .. versionadded:: 3.4
 
 
-.. function:: dump(value, fp, *, fmt=FMT_XML, sort_keys=True, skipkeys=False)
+.. function:: dump(value, fp, *, fmt=FMT_XML, sort_keys=True, skipkeys=False, aware_datetime=False)
 
    Write *value* to a plist file. *Fp* should be a writable, binary
    file object.
@@ -107,6 +114,10 @@ This module defines the following functions:
    When *skipkeys* is false (the default) the function raises :exc:`TypeError`
    when a key of a dictionary is not a string, otherwise such keys are skipped.
 
+   When *aware_datetime* is true and any field with type ``datetime.datetime``
+   is set as a :ref:`aware object <datetime-naive-aware>`, it will convert to
+   UTC timezone before writing it.
+
    A :exc:`TypeError` will be raised if the object is of an unsupported type or
    a container that contains objects of unsupported types.
 
@@ -115,8 +126,11 @@ This module defines the following functions:
 
    .. versionadded:: 3.4
 
+   .. versionchanged:: 3.13
+      The keyword-only parameter *aware_datetime* has been added.
+
 
-.. function:: dumps(value, *, fmt=FMT_XML, sort_keys=True, skipkeys=False)
+.. function:: dumps(value, *, fmt=FMT_XML, sort_keys=True, skipkeys=False, aware_datetime=False)
 
    Return *value* as a plist-formatted bytes object. See
    the documentation for :func:`dump` for an explanation of the keyword
diff --git a/Lib/plistlib.py b/Lib/plistlib.py
index 3292c30..0fc1b5c 100644
--- a/Lib/plistlib.py
+++ b/Lib/plistlib.py
@@ -140,7 +140,7 @@ def _decode_base64(s):
 _dateParser = re.compile(r"(?P<year>\d\d\d\d)(?:-(?P<month>\d\d)(?:-(?P<day>\d\d)(?:T(?P<hour>\d\d)(?::(?P<minute>\d\d)(?::(?P<second>\d\d))?)?)?)?)?Z", re.ASCII)
 
 
-def _date_from_string(s):
+def _date_from_string(s, aware_datetime):
     order = ('year', 'month', 'day', 'hour', 'minute', 'second')
     gd = _dateParser.match(s).groupdict()
     lst = []
@@ -149,10 +149,14 @@ def _date_from_string(s):
         if val is None:
             break
         lst.append(int(val))
+    if aware_datetime:
+        return datetime.datetime(*lst, tzinfo=datetime.UTC)
     return datetime.datetime(*lst)
 
 
-def _date_to_string(d):
+def _date_to_string(d, aware_datetime):
+    if aware_datetime:
+        d = d.astimezone(datetime.UTC)
     return '%04d-%02d-%02dT%02d:%02d:%02dZ' % (
         d.year, d.month, d.day,
         d.hour, d.minute, d.second
@@ -171,11 +175,12 @@ def _escape(text):
     return text
 
 class _PlistParser:
-    def __init__(self, dict_type):
+    def __init__(self, dict_type, aware_datetime=False):
         self.stack = []
         self.current_key = None
         self.root = None
         self._dict_type = dict_type
+        self._aware_datetime = aware_datetime
 
     def parse(self, fileobj):
         self.parser = ParserCreate()
@@ -277,7 +282,8 @@ class _PlistParser:
         self.add_object(_decode_base64(self.get_data()))
 
     def end_date(self):
-        self.add_object(_date_from_string(self.get_data()))
+        self.add_object(_date_from_string(self.get_data(),
+                                          aware_datetime=self._aware_datetime))
 
 
 class _DumbXMLWriter:
@@ -321,13 +327,14 @@ class _DumbXMLWriter:
 class _PlistWriter(_DumbXMLWriter):
     def __init__(
             self, file, indent_level=0, indent=b"\t", writeHeader=1,
-            sort_keys=True, skipkeys=False):
+            sort_keys=True, skipkeys=False, aware_datetime=False):
 
         if writeHeader:
             file.write(PLISTHEADER)
         _DumbXMLWriter.__init__(self, file, indent_level, indent)
         self._sort_keys = sort_keys
         self._skipkeys = skipkeys
+        self._aware_datetime = aware_datetime
 
     def write(self, value):
         self.writeln("<plist version=\"1.0\">")
@@ -360,7 +367,8 @@ class _PlistWriter(_DumbXMLWriter):
             self.write_bytes(value)
 
         elif isinstance(value, datetime.datetime):
-            self.simple_element("date", _date_to_string(value))
+            self.simple_element("date",
+                                _date_to_string(value, self._aware_datetime))
 
         elif isinstance(value, (tuple, list)):
             self.write_array(value)
@@ -461,8 +469,9 @@ class _BinaryPlistParser:
 
     see also: http://opensource.apple.com/source/CF/CF-744.18/CFBinaryPList.c
     """
-    def __init__(self, dict_type):
+    def __init__(self, dict_type, aware_datetime=False):
         self._dict_type = dict_type
+        self._aware_datime = aware_datetime
 
     def parse(self, fp):
         try:
@@ -556,8 +565,11 @@ class _BinaryPlistParser:
             f = struct.unpack('>d', self._fp.read(8))[0]
             # timestamp 0 of binary plists corresponds to 1/1/2001
             # (year of Mac OS X 10.0), instead of 1/1/1970.
-            result = (datetime.datetime(2001, 1, 1) +
-                      datetime.timedelta(seconds=f))
+            if self._aware_datime:
+                epoch = datetime.datetime(2001, 1, 1, tzinfo=datetime.UTC)
+            else:
+                epoch = datetime.datetime(2001, 1, 1)
+            result = epoch + datetime.timedelta(seconds=f)
 
         elif tokenH == 0x40:  # data
             s = self._get_size(tokenL)
@@ -629,10 +641,11 @@ def _count_to_size(count):
 _scalars = (str, int, float, datetime.datetime, bytes)
 
 class _BinaryPlistWriter (object):
-    def __init__(self, fp, sort_keys, skipkeys):
+    def __init__(self, fp, sort_keys, skipkeys, aware_datetime=False):
         self._fp = fp
         self._sort_keys = sort_keys
         self._skipkeys = skipkeys
+        self._aware_datetime = aware_datetime
 
     def write(self, value):
 
@@ -778,7 +791,12 @@ class _BinaryPlistWriter (object):
             self._fp.write(struct.pack('>Bd', 0x23, value))
 
         elif isinstance(value, datetime.datetime):
-            f = (value - datetime.datetime(2001, 1, 1)).total_seconds()
+            if self._aware_datetime:
+                dt = value.astimezone(datetime.UTC)
+                offset = dt - datetime.datetime(2001, 1, 1, tzinfo=datetime.UTC)
+                f = offset.total_seconds()
+            else:
+                f = (value - datetime.datetime(2001, 1, 1)).total_seconds()
             self._fp.write(struct.pack('>Bd', 0x33, f))
 
         elif isinstance(value, (bytes, bytearray)):
@@ -862,7 +880,7 @@ _FORMATS={
 }
 
 
-def load(fp, *, fmt=None, dict_type=dict):
+def load(fp, *, fmt=None, dict_type=dict, aware_datetime=False):
     """Read a .plist file. 'fp' should be a readable and binary file object.
     Return the unpacked root object (which usually is a dictionary).
     """
@@ -880,32 +898,36 @@ def load(fp, *, fmt=None, dict_type=dict):
     else:
         P = _FORMATS[fmt]['parser']
 
-    p = P(dict_type=dict_type)
+    p = P(dict_type=dict_type, aware_datetime=aware_datetime)
     return p.parse(fp)
 
 
-def loads(value, *, fmt=None, dict_type=dict):
+def loads(value, *, fmt=None, dict_type=dict, aware_datetime=False):
     """Read a .plist file from a bytes object.
     Return the unpacked root object (which usually is a dictionary).
     """
     fp = BytesIO(value)
-    return load(fp, fmt=fmt, dict_type=dict_type)
+    return load(fp, fmt=fmt, dict_type=dict_type, aware_datetime=aware_datetime)
 
 
-def dump(value, fp, *, fmt=FMT_XML, sort_keys=True, skipkeys=False):
+def dump(value, fp, *, fmt=FMT_XML, sort_keys=True, skipkeys=False,
+         aware_datetime=False):
     """Write 'value' to a .plist file. 'fp' should be a writable,
     binary file object.
     """
     if fmt not in _FORMATS:
         raise ValueError("Unsupported format: %r"%(fmt,))
 
-    writer = _FORMATS[fmt]["writer"](fp, sort_keys=sort_keys, skipkeys=skipkeys)
+    writer = _FORMATS[fmt]["writer"](fp, sort_keys=sort_keys, skipkeys=skipkeys,
+                                     aware_datetime=aware_datetime)
     writer.write(value)
 
 
-def dumps(value, *, fmt=FMT_XML, skipkeys=False, sort_keys=True):
+def dumps(value, *, fmt=FMT_XML, skipkeys=False, sort_keys=True,
+          aware_datetime=False):
     """Return a bytes object with the contents for a .plist file.
     """
     fp = BytesIO()
-    dump(value, fp, fmt=fmt, skipkeys=skipkeys, sort_keys=sort_keys)
+    dump(value, fp, fmt=fmt, skipkeys=skipkeys, sort_keys=sort_keys,
+         aware_datetime=aware_datetime)
     return fp.getvalue()
diff --git a/Lib/test/test_plistlib.py b/Lib/test/test_plistlib.py
index b08abab..d41975f 100644
--- a/Lib/test/test_plistlib.py
+++ b/Lib/test/test_plistlib.py
@@ -13,6 +13,8 @@ import codecs
 import subprocess
 import binascii
 import collections
+import time
+import zoneinfo
 from test import support
 from test.support import os_helper
 from io import BytesIO
@@ -838,6 +840,54 @@ class TestPlistlib(unittest.TestCase):
                                     "XML entity declarations are not supported"):
             plistlib.loads(XML_PLIST_WITH_ENTITY, fmt=plistlib.FMT_XML)
 
+    def test_load_aware_datetime(self):
+        dt = plistlib.loads(b"<plist><date>2023-12-10T08:03:30Z</date></plist>",
+                            aware_datetime=True)
+        self.assertEqual(dt.tzinfo, datetime.UTC)
+
+    @unittest.skipUnless("America/Los_Angeles" in zoneinfo.available_timezones(),
+                         "Can't find timezone datebase")
+    def test_dump_aware_datetime(self):
+        dt = datetime.datetime(2345, 6, 7, 8, 9, 10,
+                               tzinfo=zoneinfo.ZoneInfo("America/Los_Angeles"))
+        for fmt in ALL_FORMATS:
+            s = plistlib.dumps(dt, fmt=fmt, aware_datetime=True)
+            loaded_dt = plistlib.loads(s, fmt=fmt, aware_datetime=True)
+            self.assertEqual(loaded_dt.tzinfo, datetime.UTC)
+            self.assertEqual(loaded_dt, dt)
+
+    def test_dump_utc_aware_datetime(self):
+        dt = datetime.datetime(2345, 6, 7, 8, 9, 10, tzinfo=datetime.UTC)
+        for fmt in ALL_FORMATS:
+            s = plistlib.dumps(dt, fmt=fmt, aware_datetime=True)
+            loaded_dt = plistlib.loads(s, fmt=fmt, aware_datetime=True)
+            self.assertEqual(loaded_dt.tzinfo, datetime.UTC)
+            self.assertEqual(loaded_dt, dt)
+
+    @unittest.skipUnless("America/Los_Angeles" in zoneinfo.available_timezones(),
+                         "Can't find timezone datebase")
+    def test_dump_aware_datetime_without_aware_datetime_option(self):
+        dt = datetime.datetime(2345, 6, 7, 8,
+                               tzinfo=zoneinfo.ZoneInfo("America/Los_Angeles"))
+        s = plistlib.dumps(dt, fmt=plistlib.FMT_XML, aware_datetime=False)
+        self.assertIn(b"2345-06-07T08:00:00Z", s)
+
+    def test_dump_utc_aware_datetime_without_aware_datetime_option(self):
+        dt = datetime.datetime(2345, 6, 7, 8, tzinfo=datetime.UTC)
+        s = plistlib.dumps(dt, fmt=plistlib.FMT_XML, aware_datetime=False)
+        self.assertIn(b"2345-06-07T08:00:00Z", s)
+
+    def test_dump_naive_datetime_with_aware_datetime_option(self):
+        # Save a naive datetime with aware_datetime set to true.  This will lead
+        # to having different time as compared to the current machine's
+        # timezone, which is UTC.
+        dt = datetime.datetime(2345, 6, 7, 8, tzinfo=None)
+        for fmt in ALL_FORMATS:
+            s = plistlib.dumps(dt, fmt=fmt, aware_datetime=True)
+            parsed = plistlib.loads(s, aware_datetime=False)
+            expected = dt + datetime.timedelta(seconds=time.timezone)
+            self.assertEqual(parsed, expected)
+
 
 class TestBinaryPlistlib(unittest.TestCase):
 
@@ -962,6 +1012,28 @@ class TestBinaryPlistlib(unittest.TestCase):
                 with self.assertRaises(plistlib.InvalidFileException):
                     plistlib.loads(b'bplist00' + data, fmt=plistlib.FMT_BINARY)
 
+    def test_load_aware_datetime(self):
+        data = (b'bplist003B\x04>\xd0d\x00\x00\x00\x08\x00\x00\x00\x00\x00\x00'
+                b'\x01\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00'
+                b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x11')
+        self.assertEqual(plistlib.loads(data, aware_datetime=True),
+                         datetime.datetime(2345, 6, 7, 8, tzinfo=datetime.UTC))
+
+    @unittest.skipUnless("America/Los_Angeles" in zoneinfo.available_timezones(),
+                         "Can't find timezone datebase")
+    def test_dump_aware_datetime_without_aware_datetime_option(self):
+        dt = datetime.datetime(2345, 6, 7, 8,
+                               tzinfo=zoneinfo.ZoneInfo("America/Los_Angeles"))
+        msg = "can't subtract offset-naive and offset-aware datetimes"
+        with self.assertRaisesRegex(TypeError, msg):
+            plistlib.dumps(dt, fmt=plistlib.FMT_BINARY, aware_datetime=False)
+
+    def test_dump_utc_aware_datetime_without_aware_datetime_option(self):
+        dt = datetime.datetime(2345, 6, 7, 8, tzinfo=datetime.UTC)
+        msg = "can't subtract offset-naive and offset-aware datetimes"
+        with self.assertRaisesRegex(TypeError, msg):
+            plistlib.dumps(dt, fmt=plistlib.FMT_BINARY, aware_datetime=False)
+
 
 class TestKeyedArchive(unittest.TestCase):
     def test_keyed_archive_data(self):
@@ -1072,5 +1144,6 @@ class TestPlutil(unittest.TestCase):
             self.assertEqual(p.get("HexType"), 16777228)
             self.assertEqual(p.get("IntType"), 83)
 
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/Misc/NEWS.d/next/Library/2023-12-21-23-47-42.gh-issue-53502.dercJI.rst b/Misc/NEWS.d/next/Library/2023-12-21-23-47-42.gh-issue-53502.dercJI.rst
new file mode 100644
index 0000000..aa72741
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2023-12-21-23-47-42.gh-issue-53502.dercJI.rst
@@ -0,0 +1,2 @@
+Add a new option ``aware_datetime`` in :mod:`plistlib` to loads or dumps
+aware datetime.