Allow tzinfo.fromutc() to accept naive datetime instances and ensure consistent behavior

3 files changed, 60 insertions, 1 deletions

diff --git a/src/pytz/__init__.py b/src/pytz/__init__.py
index 936a072..a8c9538 100644
--- a/src/pytz/__init__.py
+++ b/src/pytz/__init__.py
@@ -212,6 +212,11 @@ class UTC(datetime.tzinfo):
     _dst = ZERO
     _tzname = zone
 
+    def fromutc(self, dt):
+        if dt.tzinfo is None:
+            return self.localize(dt)
+        return super(utc.__class__, self).fromutc(dt)
+
     def utcoffset(self, dt):
         return ZERO
 
diff --git a/src/pytz/tests/test_tzinfo.py b/src/pytz/tests/test_tzinfo.py
index dbddde5..1543fa8 100644
--- a/src/pytz/tests/test_tzinfo.py
+++ b/src/pytz/tests/test_tzinfo.py
@@ -17,7 +17,7 @@ if __name__ == '__main__':
 import pytz
 from pytz import reference
 from pytz.tzfile import _byte_string
-from pytz.tzinfo import StaticTzInfo
+from pytz.tzinfo import DstTzInfo, StaticTzInfo
 
 # I test for expected version to ensure the correct version of pytz is
 # actually being tested.
@@ -701,6 +701,56 @@ class CommonTimezonesTestCase(unittest.TestCase):
         self.assertFalse('Europe/Belfast' in pytz.common_timezones_set)
 
 
+class BaseTzInfoTestCase:
+    '''Ensure UTC, StaticTzInfo and DstTzInfo work consistently.
+
+    These tests are run for each type of tzinfo.
+    '''
+    tz = None  # override
+    tzclass = None  # override
+
+    def test_expectedclass(self):
+        self.assertTrue(isinstance(self.tz, self.tz_class))
+
+    def test_fromutc(self):
+        # naive datetime.
+        dt1 = datetime(2011, 10, 31)
+
+        # localized datetime, same timezone.
+        dt2 = self.tz.localize(dt1)
+
+        # Both should give the same results. Note that the standard
+        # Python tzinfo.fromutc() only supports the second.
+        for dt in [dt1, dt2]:
+            loc_dt = self.tz.fromutc(dt)
+            loc_dt2 = pytz.utc.localize(dt1).astimezone(self.tz)
+            self.assertEqual(loc_dt, loc_dt2)
+
+        # localized datetime, different timezone.
+        new_tz = pytz.timezone('Europe/Paris')
+        self.assertTrue(self.tz is not new_tz)
+        dt3 = new_tz.localize(dt1)
+        self.assertRaises(ValueError, self.tz.fromutc, dt3)
+
+
+class UTCTestCase(unittest.TestCase, BaseTzInfoTestCase):
+    tz = pytz.utc
+    tz_class = tz.__class__
+
+
+class StaticTzInfoTestCase(unittest.TestCase, BaseTzInfoTestCase):
+    tz = pytz.timezone('GMT')
+    tz_class = StaticTzInfo
+
+
+class DstTzInfoTestCase(unittest.TestCase, BaseTzInfoTestCase):
+    tz = pytz.timezone('Australia/Melbourne')
+    tz_class = DstTzInfo
+
+    def test_isDstTzInfo(self):
+        self.assertTrue(isinstance(self.tz, DstTzInfo))
+
+
 def test_suite():
     suite = unittest.TestSuite()
     suite.addTest(doctest.DocTestSuite('pytz'))
diff --git a/src/pytz/tzinfo.py b/src/pytz/tzinfo.py
index 1bb57c2..b61bb39 100644
--- a/src/pytz/tzinfo.py
+++ b/src/pytz/tzinfo.py
@@ -74,6 +74,8 @@ class StaticTzInfo(BaseTzInfo):
     '''
     def fromutc(self, dt):
         '''See datetime.tzinfo.fromutc'''
+        if dt.tzinfo is not None and dt.tzinfo is not self:
+            raise ValueError('fromutc: dt.tzinfo is not self')
         return (dt + self._utcoffset).replace(tzinfo=self)
 
     def utcoffset(self, dt, is_dst=None):
@@ -176,6 +178,8 @@ class DstTzInfo(BaseTzInfo):
 
     def fromutc(self, dt):
         '''See datetime.tzinfo.fromutc'''
+        if dt.tzinfo is not None and dt.tzinfo._tzinfos is not self._tzinfos:
+            raise ValueError('fromutc: dt.tzinfo is not self')
         dt = dt.replace(tzinfo=None)
         idx = max(0, bisect_right(self._utc_transition_times, dt) - 1)
         inf = self._transition_info[idx]