diff --git a/freezegun/api.py b/freezegun/api.py index 8f8c7fc6..d7ec5f5c 100644 --- a/freezegun/api.py +++ b/freezegun/api.py @@ -236,14 +236,19 @@ def start(self): continue if mod_name.startswith(tuple(self.ignore)): continue - if hasattr(module, "__name__") and module.__name__ != 'datetime': - if hasattr(module, 'datetime') and module.datetime == real_datetime: - module.datetime = FakeDatetime - if hasattr(module, 'date') and module.date == real_date: - module.date = FakeDate - if hasattr(module, "__name__") and module.__name__ != 'time': - if hasattr(module, 'time') and module.time == real_time: - module.time = fake_time + if (not hasattr(module, "__name__") + or module.__name__ in ['datetime', 'time']): + continue + for module_attribute in dir(module): + if module_attribute in ['real_date', 'real_datetime', 'real_time']: + continue + attribute_value = getattr(module, module_attribute) + if attribute_value == real_datetime: + setattr(module, module_attribute, FakeDatetime) + if attribute_value == real_date: + setattr(module, module_attribute, FakeDate) + if attribute_value == real_time: + setattr(module, module_attribute, fake_time) datetime.datetime.times_to_freeze.append(self.time_to_freeze) datetime.datetime.tz_offsets.append(self.tz_offset) @@ -271,14 +276,18 @@ def stop(self): continue if mod_name.startswith(tuple(self.ignore)): continue - if mod_name != 'datetime': - if hasattr(module, 'datetime') and module.datetime == FakeDatetime: - module.datetime = real_datetime - if hasattr(module, 'date') and module.date == FakeDate: - module.date = real_date - if mod_name != 'time': - if hasattr(module, 'time') and isinstance(module.time, FakeTime): - module.time = real_time + if mod_name in ['datetime', 'time']: + continue + for module_attribute in dir(module): + if module_attribute in ['FakeDate', 'FakeDatetime', 'FakeTime']: + continue + attribute_value = getattr(module, module_attribute) + if attribute_value == FakeDatetime: + setattr(module, module_attribute, real_datetime) + if attribute_value == FakeDate: + setattr(module, module_attribute, real_date) + if isinstance(attribute_value, FakeTime): + setattr(module, module_attribute, real_time) def decorate_callable(self, func): def wrapper(*args, **kwargs): diff --git a/tests/test_datetimes.py b/tests/test_datetimes.py index 68074fc5..8865da19 100644 --- a/tests/test_datetimes.py +++ b/tests/test_datetimes.py @@ -7,7 +7,7 @@ from tests import utils from freezegun import freeze_time -from freezegun.api import FakeDatetime, FakeDate, real_date +from freezegun.api import FakeDatetime, FakeDate class temp_locale(object): @@ -259,7 +259,8 @@ def assert_class_of_datetimes(right_class, wrong_class): def test_min_and_max(): freezer = freeze_time("2012-01-14") - real_datetime = datetime + real_datetime = datetime.datetime + real_date = datetime.date freezer.start() datetime.datetime.min.__class__.should.equal(FakeDatetime) diff --git a/tests/test_import_alias.py b/tests/test_import_alias.py new file mode 100644 index 00000000..6e40d924 --- /dev/null +++ b/tests/test_import_alias.py @@ -0,0 +1,13 @@ +from freezegun import freeze_time +from datetime import datetime as datetime_aliased +from time import time as time_aliased + + +@freeze_time("1980-01-01") +def test_datetime_alias(): + assert datetime_aliased.now() == datetime_aliased(1980,1,1) + + +@freeze_time("1970-01-01") +def test_time_alias(): + assert time_aliased() == 0.0