From e7d5392f6e5472122b3b6035ce2c6f70a9440d19 Mon Sep 17 00:00:00 2001
From: Bolke de Bruin <bolke@xs4all.nl>
Date: Wed, 15 Nov 2017 00:03:47 +0100
Subject: [PATCH] Add deepcopy / pickle compatibility

Python's doc say: "A tzinfo subclass must have an __init__() method that can
be called with no arguments". Timezone and TimezoneInfo don't honor this
requirement. Defining __getinitargs__ is sufficient to fix copy/deepcopy as
well as pickling/unpickling.

This introduces __tzinfos to Timezone as otherwise recursion errors will happen
due to TimezoneInfo having a reference to Timezone.
---
 pendulum/tz/timezone.py               | 15 +++++++++++++++
 pendulum/tz/timezone_info.py          |  2 ++
 tests/pendulum_tests/test_behavior.py | 25 +++++++++++++++++++++++++
 3 files changed, 42 insertions(+)

diff --git a/pendulum/tz/timezone.py b/pendulum/tz/timezone.py
index 33cd0428..88592396 100644
--- a/pendulum/tz/timezone.py
+++ b/pendulum/tz/timezone.py
@@ -48,6 +48,7 @@ def __init__(self, name, transitions,
         """
         self._name = name
         self._transitions = transitions
+        self.__tzinfos = tzinfos
         self._tzinfos = tuple(
             map(lambda tzinfo: TimezoneInfo(self, *tzinfo), tzinfos)
         )
@@ -449,6 +450,11 @@ def _find_utc_index(self, dt):
     def __repr__(self):
         return '<Timezone [{}]>'.format(self._name)
 
+    def __getinitargs__(self):
+        return (self._name, self._transitions,
+                self.__tzinfos, self._default_tzinfo_index,
+                self._utc_transition_times)
+
 
 class FixedTimezone(Timezone):
     """
@@ -486,6 +492,8 @@ def __init__(self, offset, name=None, transition_type=None):
         )
         self._tzinfo = self._tzinfos[0]
 
+        self._offset = offset
+
     @classmethod
     def load(cls, name):
         if name not in cls._cache:
@@ -513,6 +521,9 @@ def fromutc(self, dt):
 
         return (dt + self._tzinfo.adjusted_offset).replace(tzinfo=self._tzinfo)
 
+    def __getinitargs__(self):
+        return self._offset
+
 
 class _UTC(FixedTimezone):
 
@@ -524,4 +535,8 @@ def __init__(self):
     def fromutc(self, dt):
         return dt.replace(tzinfo=UTC)
 
+    def __getinitargs__(self):
+        return ()
+
+
 UTCTimezone = _UTC()
diff --git a/pendulum/tz/timezone_info.py b/pendulum/tz/timezone_info.py
index 51d5afaf..8dc1f4a7 100644
--- a/pendulum/tz/timezone_info.py
+++ b/pendulum/tz/timezone_info.py
@@ -113,6 +113,8 @@ def __repr__(self):
             'DST' if self.is_dst else 'STD',
         )
 
+    def __getinitargs__(self):
+        return self._tz, self._utc_offset, self._is_dst, self._dst, self._abbrev
 
 class _UTC(TimezoneInfo):
 
diff --git a/tests/pendulum_tests/test_behavior.py b/tests/pendulum_tests/test_behavior.py
index ce8a8173..12a12854 100644
--- a/tests/pendulum_tests/test_behavior.py
+++ b/tests/pendulum_tests/test_behavior.py
@@ -2,8 +2,10 @@
 
 import pickle
 import pendulum
+from copy import deepcopy
 from datetime import datetime, date, time, timedelta
 from pendulum import Pendulum, timezone
+from pendulum.tz.timezone import Timezone
 from .. import AbstractTestCase
 
 
@@ -103,3 +105,26 @@ def test_proper_dst(self):
         dt = pendulum.create(1941, 7, 1, tz='Europe/Amsterdam')
 
         self.assertEqual(timedelta(0, 6000), dt.dst())
+
+    def test_deepcopy(self):
+        dt = pendulum.create(1941, 7, 1, tz='Europe/Amsterdam')
+
+        self.assertEqual(dt, deepcopy(dt))
+
+    def test_deepcopy_datetime(self):
+        dt = pendulum.create(1941, 7, 1, tz='Europe/Amsterdam')
+
+        self.assertEqual(dt._datetime, deepcopy(dt._datetime))
+
+    def test_pickle_timezone(self):
+        dt1 = pendulum.timezone('Europe/Amsterdam')
+        s = pickle.dumps(dt1)
+        dt2 = pickle.loads(s)
+
+        self.assertTrue(isinstance(dt2, Timezone))
+
+        dt1 = pendulum.timezone('UTC')
+        s = pickle.dumps(dt1)
+        dt2 = pickle.loads(s)
+
+        self.assertTrue(isinstance(dt2, Timezone))