Repository: joke2k/faker
Pull request #299 (2015-11-09 12:58:41)

Add tzinfo support to datetime methods


Patch:
diff --git a/faker/providers/date_time/__init__.py b/faker/providers/date_time/__init__.py
index f2e93e5c..06637d7b 100644
--- a/faker/providers/date_time/__init__.py
+++ b/faker/providers/date_time/__init__.py
@@ -242,18 +242,20 @@ class Provider(BaseProvider):
         return timedelta(seconds=ts)
 
     @classmethod
-    def date_time(cls):
+    def date_time(cls, tzinfo=None):
         """
         Get a datetime object for a date between January 1, 1970 and now
+        :param tzinfo: timezone, instance of datetime.tzinfo subclass
         :example DateTime('2005-08-16 20:39:21')
         :return datetime
         """
-        return datetime.fromtimestamp(cls.unix_time())
+        return datetime.fromtimestamp(cls.unix_time(), tzinfo)
 
     @classmethod
-    def date_time_ad(cls):
+    def date_time_ad(cls, tzinfo=None):
         """
         Get a datetime object for a date between January 1, 001 and now
+        :param tzinfo: timezone, instance of datetime.tzinfo subclass
         :example DateTime('1265-03-22 21:15:52')
         :return datetime
         """
@@ -262,14 +264,15 @@ class Provider(BaseProvider):
         #       a "ValueError: timestamp out of range for platform time_t"
         #       on some platforms due to system C functions;
         #       see http://stackoverflow.com/a/10588133/2315612
-        return datetime.fromtimestamp(0) + timedelta(seconds=ts)
+        return datetime.fromtimestamp(0, tzinfo) + timedelta(seconds=ts)
 
     @classmethod
-    def iso8601(cls):
+    def iso8601(cls, tzinfo=None):
         """
+        :param tzinfo: timezone, instance of datetime.tzinfo subclass
         :example '2003-10-21T16:05:52+0000'
         """
-        return cls.date_time().isoformat()
+        return cls.date_time(tzinfo).isoformat()
 
     @classmethod
     def date(cls, pattern='%Y-%m-%d'):
@@ -290,15 +293,15 @@ class Provider(BaseProvider):
         return cls.date_time().time().strftime(pattern)
 
     @classmethod
-    def _parse_date_time(cls, text):
+    def _parse_date_time(cls, text, tzinfo=None):
         if isinstance(text, (datetime, date, real_datetime, real_date)):
             return datetime_to_timestamp(text)
-        now = datetime.now()
+        now = datetime.now(tzinfo)
         if isinstance(text, timedelta):
             return datetime_to_timestamp(now - text)
         if is_string(text):
             if text == 'now':
-                return datetime_to_timestamp(datetime.now())
+                return datetime_to_timestamp(datetime.now(tzinfo))
             parts = cls.regex.match(text)
             if not parts:
                 return
@@ -319,134 +322,140 @@ class Provider(BaseProvider):
         raise ValueError("Invalid format for date '{0}'".format(text))
 
     @classmethod
-    def date_time_between(cls, start_date='-30y', end_date='now'):
+    def date_time_between(cls, start_date='-30y', end_date='now', tzinfo=None):
         """
         Get a DateTime object based on a random date between two given dates.
         Accepts date strings that can be recognized by strtotime().
 
         :param start_date Defaults to 30 years ago
         :param end_date Defaults to "now"
+        :param tzinfo: timezone, instance of datetime.tzinfo subclass
         :example DateTime('1999-02-02 11:42:52')
         :return DateTime
         """
         start_date = cls._parse_date_time(start_date)
         end_date = cls._parse_date_time(end_date)
         timestamp = random.randint(start_date, end_date)
-        return datetime.fromtimestamp(timestamp)
+        return datetime.fromtimestamp(timestamp, tzinfo)
 
     @classmethod
-    def date_time_between_dates(cls, datetime_start=None, datetime_end=None):
+    def date_time_between_dates(cls, datetime_start=None, datetime_end=None, tzinfo=None):
         """
         Takes two DateTime objects and returns a random date between the two given dates.
         Accepts DateTime objects.
 
         :param datetime_start DateTime
         :param datetime_end DateTime
+        :param tzinfo: timezone, instance of datetime.tzinfo subclass
         :example DateTime('1999-02-02 11:42:52')
         :return DateTime
         """
         if datetime_start is None:
-            datetime_start = datetime.now()
+            datetime_start = datetime.now(tzinfo)
 
         if datetime_end is None:
-            datetime_end = datetime.now()
+            datetime_end = datetime.now(tzinfo)
 
         timestamp = random.randint(datetime_to_timestamp(datetime_start),
                                        datetime_to_timestamp(datetime_end))
-        return datetime.fromtimestamp(timestamp)
+        return datetime.fromtimestamp(timestamp, tzinfo)
 
     @classmethod
-    def date_time_this_century(cls, before_now=True, after_now=False):
+    def date_time_this_century(cls, before_now=True, after_now=False, tzinfo=None):
         """
-        Gets a DateTime object for the decade year. 
-        
+        Gets a DateTime object for the decade year.
+
         :param before_now: include days in current decade before today
         :param after_now: include days in current decade after today
+        :param tzinfo: timezone, instance of datetime.tzinfo subclass
         :example DateTime('2012-04-04 11:02:02')
         :return DateTime
         """
-        now = datetime.now()
-        this_century_start = datetime(now.year - (now.year % 10), 1, 1)
-        next_century_start = datetime(this_century_start.year + 10, 1, 1)
+        now = datetime.now(tzinfo)
+        this_century_start = datetime(now.year - (now.year % 10), 1, 1, tzinfo=tzinfo)
+        next_century_start = datetime(this_century_start.year + 10, 1, 1, tzinfo=tzinfo)
 
         if before_now and after_now:
-            return cls.date_time_between_dates(this_century_start, next_century_start)
+            return cls.date_time_between_dates(this_century_start, next_century_start, tzinfo)
         elif not before_now and after_now:
-            return cls.date_time_between_dates(now, next_century_start)
+            return cls.date_time_between_dates(now, next_century_start, tzinfo)
         elif not after_now and before_now:
-            return cls.date_time_between_dates(this_century_start, now)
+            return cls.date_time_between_dates(this_century_start, now, tzinfo)
         else:
             return now
 
     @classmethod
-    def date_time_this_decade(cls, before_now=True, after_now=False):
+    def date_time_this_decade(cls, before_now=True, after_now=False, tzinfo=None):
         """
-        Gets a DateTime object for the decade year. 
-        
+        Gets a DateTime object for the decade year.
+
         :param before_now: include days in current decade before today
         :param after_now: include days in current decade after today
+        :param tzinfo: timezone, instance of datetime.tzinfo subclass
         :example DateTime('2012-04-04 11:02:02')
         :return DateTime
         """
-        now = datetime.now()
-        this_decade_start = datetime(now.year - (now.year % 10), 1, 1)
-        next_decade_start = datetime(this_decade_start.year + 10, 1, 1)
+        now = datetime.now(tzinfo)
+        this_decade_start = datetime(now.year - (now.year % 10), 1, 1, tzinfo=tzinfo)
+        next_decade_start = datetime(this_decade_start.year + 10, 1, 1, tzinfo=tzinfo)
 
         if before_now and after_now:
-            return cls.date_time_between_dates(this_decade_start, next_decade_start)
+            return cls.date_time_between_dates(this_decade_start, next_decade_start, tzinfo)
         elif not before_now and after_now:
-            return cls.date_time_between_dates(now, next_decade_start)
+            return cls.date_time_between_dates(now, next_decade_start, tzinfo)
         elif not after_now and before_now:
-            return cls.date_time_between_dates(this_decade_start, now)
+            return cls.date_time_between_dates(this_decade_start, now, tzinfo)
         else:
             return now
 
     @classmethod
-    def date_time_this_year(cls, before_now=True, after_now=False):
+    def date_time_this_year(cls, before_now=True, after_now=False, tzinfo=None):
         """
-        Gets a DateTime object for the current year. 
-        
+        Gets a DateTime object for the current year.
+
         :param before_now: include days in current year before today
         :param after_now: include days in current year after today
+        :param tzinfo: timezone, instance of datetime.tzinfo subclass
         :example DateTime('2012-04-04 11:02:02')
         :return DateTime
         """
-        now = datetime.now()
+        now = datetime.now(tzinfo)
         this_year_start = now.replace(month=1, day=1, hour=0, minute=0, second=0, microsecond=0)
-        next_year_start = datetime(now.year + 1, 1, 1)
+        next_year_start = datetime(now.year + 1, 1, 1, tzinfo=tzinfo)
 
         if before_now and after_now:
-            return cls.date_time_between_dates(this_year_start, next_year_start)
+            return cls.date_time_between_dates(this_year_start, next_year_start, tzinfo)
         elif not before_now and after_now:
-            return cls.date_time_between_dates(now, next_year_start)
+            return cls.date_time_between_dates(now, next_year_start, tzinfo)
         elif not after_now and before_now:
-            return cls.date_time_between_dates(this_year_start, now)
+            return cls.date_time_between_dates(this_year_start, now, tzinfo)
         else:
             return now
-        
+
     @classmethod
-    def date_time_this_month(cls, before_now=True, after_now=False):
+    def date_time_this_month(cls, before_now=True, after_now=False, tzinfo=None):
         """
-        Gets a DateTime object for the current month. 
-        
+        Gets a DateTime object for the current month.
+
         :param before_now: include days in current month before today
         :param after_now: include days in current month after today
+        :param tzinfo: timezone, instance of datetime.tzinfo subclass
         :example DateTime('2012-04-04 11:02:02')
         :return DateTime
         """
-        now = datetime.now()
+        now = datetime.now(tzinfo)
         this_month_start = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
-        next_month_start = datetime(now.year, now.month + 1 % 12, 1)
+        next_month_start = datetime(now.year, now.month + 1 % 12, 1, tzinfo=tzinfo)
 
         if before_now and after_now:
-            return cls.date_time_between_dates(this_month_start, next_month_start)
+            return cls.date_time_between_dates(this_month_start, next_month_start, tzinfo)
         elif not before_now and after_now:
-            return cls.date_time_between_dates(now, next_month_start)
+            return cls.date_time_between_dates(now, next_month_start, tzinfo)
         elif not after_now and before_now:
-            return cls.date_time_between_dates(this_month_start, now)
+            return cls.date_time_between_dates(this_month_start, now, tzinfo)
         else:
             return now
-            
+
     @classmethod
     def am_pm(cls):
         return cls.date('%p')


Test patch:
diff --git a/faker/tests/__init__.py b/faker/tests/__init__.py
index ec4e46a7..8019b007 100644
--- a/faker/tests/__init__.py
+++ b/faker/tests/__init__.py
@@ -34,6 +34,25 @@ except NameError:
 TEST_DIR = os.path.dirname(__file__)
 
 
+class UTC(datetime.tzinfo):
+    """
+    UTC implementation taken from Python's docs.
+    """
+    def __repr__(self):
+        return "<UTC>"
+
+    def utcoffset(self, dt):
+        return datetime.timedelta(0)
+
+    def tzname(self, dt):
+        return "UTC"
+
+    def dst(self, dt):
+        return datetime.timedelta(0)
+
+utc = UTC()
+
+
 class BarProvider(object):
     def foo_formatter(self):
         return 'barfoo'
@@ -306,7 +325,20 @@ class FactoryTestCase(unittest.TestCase):
         # test that certain formatting strings are allowed on post-1900 dates
         result = datetime_safe.date(2008, 2, 29).strftime('%y')
         self.assertEqual(result, r'08')
-    
+
+    def test_datetimes_with_and_without_tzinfo(self):
+        from faker.providers.date_time import Provider
+        provider = Provider
+
+        self.assertEqual(provider.date_time().tzinfo, None)
+        self.assertEqual(provider.date_time(utc).tzinfo, utc)
+
+        self.assertEqual(provider.date_time_ad().tzinfo, None)
+        self.assertEqual(provider.date_time_ad(utc).tzinfo, utc)
+
+        self.assertFalse(provider.iso8601().endswith('+00:00'))
+        self.assertTrue(provider.iso8601(utc).endswith('+00:00'))
+
     def test_date_time_between_dates(self):
         from faker.providers.date_time import Provider
         provider = Provider
@@ -321,6 +353,23 @@ class FactoryTestCase(unittest.TestCase):
         self.assertTrue(datetime_start <= random_date)
         self.assertTrue(datetime_end >= random_date)
 
+    def test_date_time_between_dates_with_tzinfo(self):
+        from faker.providers.date_time import Provider
+        provider = Provider
+
+        timestamp_start = random.randint(0, 10000000000)
+        timestamp_end = timestamp_start+1
+
+        datetime_start = datetime.datetime.fromtimestamp(timestamp_start, utc)
+        datetime_end = datetime.datetime.fromtimestamp(timestamp_end, utc)
+
+        random_date_naive = provider.date_time_between_dates(datetime_start, datetime_end)
+        self.assertRaises(TypeError, lambda: datetime_start <= random_date_naive)
+
+        random_date = provider.date_time_between_dates(datetime_start, datetime_end, utc)
+        self.assertTrue(datetime_start <= random_date)
+        self.assertTrue(datetime_end >= random_date)
+
     def _datetime_to_time(self, value):
         return int(time.mktime(value.timetuple()))
 
@@ -346,6 +395,36 @@ class FactoryTestCase(unittest.TestCase):
         self.assertEqual(self._datetime_to_time(provider.date_time_this_month(before_now=False, after_now=False)),
                                self._datetime_to_time(datetime.datetime.now()))
 
+    def test_date_time_this_period_with_tzinfo(self):
+        from faker.providers.date_time import Provider
+        provider = Provider
+
+        # ensure all methods provide timezone aware datetimes
+        self.assertRaises(TypeError, lambda: provider.date_time_this_century(before_now=False, after_now=True, tzinfo=utc) >= datetime.datetime.now())
+        self.assertRaises(TypeError, lambda: provider.date_time_this_decade(after_now=False, tzinfo=utc) <= datetime.datetime.now())
+        self.assertRaises(TypeError, lambda: provider.date_time_this_year(after_now=False, tzinfo=utc) <= datetime.datetime.now())
+        self.assertRaises(TypeError, lambda: provider.date_time_this_month(after_now=False, tzinfo=utc) <= datetime.datetime.now())
+
+        # test century
+        self.assertTrue(provider.date_time_this_century(after_now=False, tzinfo=utc) <= datetime.datetime.now(utc))
+        self.assertTrue(provider.date_time_this_century(before_now=False, after_now=True, tzinfo=utc) >= datetime.datetime.now(utc))
+        # test decade
+        self.assertTrue(provider.date_time_this_decade(after_now=False, tzinfo=utc) <= datetime.datetime.now(utc))
+        self.assertTrue(provider.date_time_this_decade(before_now=False, after_now=True, tzinfo=utc) >= datetime.datetime.now(utc))
+
+        self.assertEqual(provider.date_time_this_decade(before_now=False, after_now=False, tzinfo=utc).replace(microsecond=0),
+                               datetime.datetime.now(utc).replace(microsecond=0))
+        # test year
+        self.assertTrue(provider.date_time_this_year(after_now=False, tzinfo=utc) <= datetime.datetime.now(utc))
+        self.assertTrue(provider.date_time_this_year(before_now=False, after_now=True, tzinfo=utc) >= datetime.datetime.now(utc))
+        self.assertEqual(provider.date_time_this_year(before_now=False, after_now=False, tzinfo=utc).replace(microsecond=0),
+                               datetime.datetime.now(utc).replace(microsecond=0))
+        # test month
+        self.assertTrue(provider.date_time_this_month(after_now=False, tzinfo=utc) <= datetime.datetime.now(utc))
+        self.assertTrue(provider.date_time_this_month(before_now=False, after_now=True, tzinfo=utc) >= datetime.datetime.now(utc))
+        self.assertEqual(provider.date_time_this_month(before_now=False, after_now=False, tzinfo=utc).replace(microsecond=0),
+                               datetime.datetime.now(utc).replace(microsecond=0))
+
     def test_prefix_suffix_always_string(self):
         # Locales known to contain `*_male` and `*_female`.
         for locale in ("bg_BG", "dk_DK", "en", "ru_RU", "tr_TR"):