Add move_to() method (#39)

adamchainz · Jul 8, 2020 · bff3e9b · bff3e9b
1 parent 7004e16
commit bff3e9b
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 19 deletions.
diff --git a/src/time_machine.py b/src/time_machine.py
@@ -21,6 +21,30 @@
 NANOSECONDS_PER_SECOND = 1_000_000_000
 
 
+def destination_to_timestamp(destination):
+    if callable(destination):
+        destination = destination()
+    elif isinstance(destination, GeneratorType):
+        destination = next(destination)
+
+    if isinstance(destination, (int, float)):
+        destination_timestamp = destination
+    elif isinstance(destination, dt.datetime):
+        if destination.tzinfo is None:
+            destination = destination.replace(tzinfo=dt.timezone.utc)
+        destination_timestamp = destination.timestamp()
+    elif isinstance(destination, dt.date):
+        destination_timestamp = dt.datetime.combine(
+            destination, dt.time(0, 0), tzinfo=dt.timezone.utc
+        ).timestamp()
+    elif isinstance(destination, str):
+        destination_timestamp = parse_datetime(destination).timestamp()
+    else:
+        raise TypeError(f"Unsupported destination {destination!r}")
+
+    return destination_timestamp
+
+
 class Coordinates:
     def __init__(self, destination_timestamp: float, tick: bool):
         self.destination_timestamp = destination_timestamp
@@ -50,6 +74,9 @@ def shift(self, delta):
 
         self.destination_timestamp += total_seconds
 
+    def move_to(self, destination):
+        self.destination_timestamp = destination_to_timestamp(destination)
+
 
 coordinates_stack = []
 
@@ -75,25 +102,7 @@ def uuid_idempotent_load_system_functions():
 
 class travel:
     def __init__(self, destination, *, tick=True, tz_offset=None):
-        if callable(destination):
-            destination = destination()
-        elif isinstance(destination, GeneratorType):
-            destination = next(destination)
-
-        if isinstance(destination, (int, float)):
-            destination_timestamp = destination
-        elif isinstance(destination, dt.datetime):
-            if destination.tzinfo is None:
-                destination = destination.replace(tzinfo=dt.timezone.utc)
-            destination_timestamp = destination.timestamp()
-        elif isinstance(destination, dt.date):
-            destination_timestamp = dt.datetime.combine(
-                destination, dt.time(0, 0), tzinfo=dt.timezone.utc
-            ).timestamp()
-        elif isinstance(destination, str):
-            destination_timestamp = parse_datetime(destination).timestamp()
-        else:
-            raise TypeError(f"Unsupported destination {destination!r}")
+        destination_timestamp = destination_to_timestamp(destination)
 
         if tz_offset is not None:
             if isinstance(tz_offset, dt.timedelta):

diff --git a/tests/test_time_machine.py b/tests/test_time_machine.py
@@ -456,6 +456,35 @@ def test_shift_when_tick():
         assert EPOCH + 10.0 <= time.time() < EPOCH + 20.0
 
 
+# move_to() tests
+
+
+def test_move_to_datetime():
+    with time_machine.travel(EPOCH, tick=False) as traveller:
+        assert time.time() == EPOCH
+        traveller.move_to(EPOCH_PLUS_ONE_YEAR_DATETIME)
+        assert time.time() == EPOCH_PLUS_ONE_YEAR
+
+
+def test_move_to_datetime_when_tick():
+    with time_machine.travel(EPOCH, tick=True) as traveller:
+        traveller.move_to(EPOCH_PLUS_ONE_YEAR_DATETIME)
+        assert time.time() == pytest.approx(EPOCH_PLUS_ONE_YEAR)
+
+
+def test_move_to_past_datetime():
+    with time_machine.travel(EPOCH_PLUS_ONE_YEAR, tick=False) as traveller:
+        assert time.time() == EPOCH_PLUS_ONE_YEAR
+        traveller.move_to(EPOCH_DATETIME)
+        assert time.time() == EPOCH
+
+
+def test_move_to_past_datetime_when_tick():
+    with time_machine.travel(EPOCH_PLUS_ONE_YEAR_DATETIME, tick=True) as traveller:
+        traveller.move_to(EPOCH)
+        assert time.time() == pytest.approx(EPOCH)
+
+
 # uuid tests