Improved World method type signatures.

epymorph/engine/mm_exec.py

@@ -96,24 +96,22 @@ def apply(self, world: World, tick: Tick) -> None:
             total += travelers.sum()
 
         # Process return clause.
-        # return_requested, return_actual, return_total = world.apply_return(tick)
         return_movers = world.apply_return(tick, return_stats=True)
-        if return_movers is not None:
-            return_total = return_movers.sum()
-            total += return_total
-
-            self.on_movement_clause.publish(
-                OnMovementClause(
-                    tick.index,
-                    tick.day,
-                    tick.step,
-                    "return",
-                    return_movers,
-                    return_movers,
-                    return_total,
-                    False,
-                )
+        return_total = return_movers.sum()
+        total += return_total
+
+        self.on_movement_clause.publish(
+            OnMovementClause(
+                tick.index,
+                tick.day,
+                tick.step,
+                "return",
+                return_movers,
+                return_movers,
+                return_total,
+                False,
             )
+        )
 
         self.on_movement_finish.publish(
             OnMovementFinish(tick.index, tick.day, tick.step, total))

epymorph/engine/world.py

@@ -6,6 +6,7 @@
 and will eventually be returning or moving to another location.
 """
 from abc import ABC, abstractmethod
+from typing import Literal, overload
 
 from numpy.typing import NDArray
 
@@ -45,6 +46,16 @@ def apply_travel(self, travelers: NDArray[SimDType], return_tick: int) -> None:
         modify the world state as a result of that movement.
         """
 
+    @abstractmethod
+    @overload
+    def apply_return(self, tick: Tick, *, return_stats: Literal[False]) -> None:
+        ...
+
+    @abstractmethod
+    @overload
+    def apply_return(self, tick: Tick, *, return_stats: Literal[True]) -> NDArray[SimDType]:
+        ...
+
     @abstractmethod
     def apply_return(self, tick: Tick, *, return_stats: bool) -> NDArray[SimDType] | None:
         """

epymorph/engine/world_hypercube.py

@@ -1,5 +1,5 @@
 """World implementation: HypercubeWorld."""
-from typing import Self
+from typing import Literal, Self, overload
 
 import numpy as np
 import psutil
@@ -107,6 +107,14 @@ def apply_travel(self, travelers: NDArray[SimDType], return_tick: int) -> None:
         self.ledger[return_tick + 1, :, :, :] += travelers
         self.time_frontier = max(self.time_frontier, return_tick + 2)
 
+    @overload
+    def apply_return(self, tick: Tick, *, return_stats: Literal[False]) -> None:
+        ...
+
+    @overload
+    def apply_return(self, tick: Tick, *, return_stats: Literal[True]) -> NDArray[SimDType]:
+        ...
+
     def apply_return(self, tick: Tick, *, return_stats: bool) -> NDArray[SimDType] | None:
         # we have to transpose the movers "stats" result since they're being stored here as
         # (home, visiting) and our result needs to be

epymorph/engine/world_list.py

@@ -1,6 +1,6 @@
 """World implementation: ListWorld."""
 from operator import attrgetter
-from typing import Any, Iterable, Self
+from typing import Any, Iterable, Literal, Self, overload
 
 import numpy as np
 from numpy.typing import NDArray
@@ -135,6 +135,14 @@ def apply_travel(self, travelers: NDArray[SimDType], return_tick: int) -> None:
 
         self.normalize()
 
+    @overload
+    def apply_return(self, tick: Tick, *, return_stats: Literal[False]) -> None:
+        ...
+
+    @overload
+    def apply_return(self, tick: Tick, *, return_stats: Literal[True]) -> NDArray[SimDType]:
+        ...
+
     def apply_return(self, tick: Tick, *, return_stats: bool) -> NDArray[SimDType] | None:
         movers: Any = None
         if return_stats: