Improve more typing

Author: bryanforbes

pendulum/date.py

@@ -8,7 +8,6 @@
 from datetime import timedelta
 from typing import TYPE_CHECKING
 from typing import NoReturn
-from typing import TypeVar
 from typing import cast
 from typing import overload
 
@@ -33,9 +32,6 @@
     from typing_extensions import Self
 
 
-_D = TypeVar("_D", bound="Date")
-
-
 class Date(FormattableMixin, date):
     # Names of days of the week
     _days = {
@@ -275,7 +271,7 @@ def __sub__(self, __dt: datetime) -> NoReturn:
         ...
 
     @overload
-    def __sub__(self: _D, __dt: _D) -> Interval:
+    def __sub__(self, __dt: Self) -> Interval:
         ...
 
     def __sub__(self, other: timedelta | date) -> Self | Interval:

pendulum/testing/traveller.py

@@ -9,15 +9,17 @@
 if TYPE_CHECKING:
     from types import TracebackType
 
+    from typing_extensions import Self
+
 
 class BaseTraveller:
     def __init__(self, datetime_class: type[DateTime] = DateTime) -> None:
         self._datetime_class: type[DateTime] = datetime_class
 
-    def freeze(self: BaseTraveller) -> BaseTraveller:
+    def freeze(self) -> Self:
         raise NotImplementedError()
 
-    def travel_back(self: BaseTraveller) -> BaseTraveller:
+    def travel_back(self) -> Self:
         raise NotImplementedError()
 
     def travel(
@@ -30,12 +32,23 @@ def travel(
         minutes: int = 0,
         seconds: int = 0,
         microseconds: int = 0,
-    ) -> BaseTraveller:
+    ) -> Self:
         raise NotImplementedError()
 
-    def travel_to(self, dt: DateTime) -> BaseTraveller:
+    def travel_to(self, dt: DateTime, *, freeze: bool = False) -> Self:
         raise NotImplementedError()
 
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType,
+    ) -> None:
+        ...
+
 
 if not PYPY:
     import time_machine
@@ -48,7 +61,7 @@ def __init__(self, datetime_class: type[DateTime] = DateTime) -> None:
             self._traveller: time_machine.travel | None = None
             self._coordinates: time_machine.Coordinates | None = None
 
-        def freeze(self) -> Traveller:
+        def freeze(self) -> Self:
             if self._started:
                 cast(time_machine.Coordinates, self._coordinates).move_to(
                     self._datetime_class.now(), tick=False
@@ -58,7 +71,7 @@ def freeze(self) -> Traveller:
 
             return self
 
-        def travel_back(self) -> Traveller:
+        def travel_back(self) -> Self:
             if not self._started:
                 return self
 
@@ -81,7 +94,7 @@ def travel(
             microseconds: int = 0,
             *,
             freeze: bool = False,
-        ) -> Traveller:
+        ) -> Self:
             self._start(freeze=freeze)
 
             cast(time_machine.Coordinates, self._coordinates).move_to(
@@ -99,7 +112,7 @@ def travel(
 
             return self
 
-        def travel_to(self, dt: DateTime, *, freeze: bool = False) -> Traveller:
+        def travel_to(self, dt: DateTime, *, freeze: bool = False) -> Self:
             self._start(freeze=freeze)
 
             cast(time_machine.Coordinates, self._coordinates).move_to(dt)
@@ -119,7 +132,7 @@ def _start(self, freeze: bool = False) -> None:
 
             self._started = True
 
-        def __enter__(self) -> Traveller:
+        def __enter__(self) -> Self:
             self._start()
 
             return self

pendulum/tz/__init__.py

@@ -2,6 +2,7 @@
 
 from pathlib import Path
 from typing import cast
+from typing import overload
 
 from pendulum.tz.local_timezone import get_local_timezone
 from pendulum.tz.local_timezone import set_local_timezone
@@ -30,6 +31,21 @@ def timezones() -> tuple[str, ...]:
     return _timezones
 
 
+@overload
+def timezone(name: int) -> FixedTimezone:
+    ...
+
+
+@overload
+def timezone(name: str) -> Timezone:
+    ...
+
+
+@overload
+def timezone(name: str | int) -> Timezone | FixedTimezone:
+    ...
+
+
 def timezone(name: str | int) -> Timezone | FixedTimezone:
     """
     Return a Timezone instance given its name.

poetry.lock

@@ -53,6 +53,7 @@ pre-commit = "^2.20.0"
 [tool.poetry.group.typing.dependencies]
 mypy = "^1.0.1"
 types-python-dateutil = "^2.8.19"
+types-pytz = ">=2022.7.1.2"
 
 [tool.poetry.group.dev.dependencies]
 babel = "^2.10.3"
@@ -111,9 +112,6 @@ exclude = [
 [[tool.mypy.overrides]]
 module = [
     "pendulum.mixins.default",
-    "tests.conftest",
-    "tests.test_helpers",
-    "tests.test_main",
     "tests.test_parsing",
     "tests.date.test_add",
     "tests.date.test_behavior",

pyproject.toml

@@ -1,12 +1,17 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import pytest
 
 import pendulum
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 
 @pytest.fixture(autouse=True)
-def setup():
+def setup() -> Iterator[None]:
     pendulum.set_local_timezone(pendulum.timezone("America/Toronto"))
 
     yield
@@ -18,8 +23,15 @@ def setup():
 
 
 def assert_datetime(
-    d, year, month, day, hour=None, minute=None, second=None, microsecond=None
-):
+    d: pendulum.DateTime,
+    year: int,
+    month: int,
+    day: int,
+    hour: int | None = None,
+    minute: int | None = None,
+    second: int | None = None,
+    microsecond: int | None = None,
+) -> None:
     assert year == d.year
     assert month == d.month
     assert day == d.day
@@ -37,13 +49,19 @@ def assert_datetime(
         assert microsecond == d.microsecond
 
 
-def assert_date(d, year, month, day):
+def assert_date(d: pendulum.Date, year: int, month: int, day: int) -> None:
     assert year == d.year
     assert month == d.month
     assert day == d.day
 
 
-def assert_time(t, hour, minute, second, microsecond=None):
+def assert_time(
+    t: pendulum.Time,
+    hour: int,
+    minute: int,
+    second: int,
+    microsecond: int | None = None,
+) -> None:
     assert hour == t.hour
     assert minute == t.minute
     assert second == t.second
@@ -53,15 +71,15 @@ def assert_time(t, hour, minute, second, microsecond=None):
 
 
 def assert_duration(
-    dur,
-    years=None,
-    months=None,
-    weeks=None,
-    days=None,
-    hours=None,
-    minutes=None,
-    seconds=None,
-    microseconds=None,
+    dur: pendulum.Duration,
+    years: int | None = None,
+    months: int | None = None,
+    weeks: int | None = None,
+    days: int | None = None,
+    hours: int | None = None,
+    minutes: int | None = None,
+    seconds: int | None = None,
+    microseconds: int | None = None,
 ) -> None:
     expected = {}
     actual = {}