Allow specifying timezone

CHANGELOG.md

@@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file.
 
 ## v0.4.0
 
+### Added
+
+- `YearMonth.current()` now accepts a `tz` argument to specify the timezone.
+
 ### Changed
 
 - Made `YearMonth` instances immutable.

README.md

@@ -17,7 +17,21 @@ ym = YearMonth(2019, 1)
 print(ym) # 2019-01
 ```
 
-Parsing ISO 8601 strings:
+### Getting the current year-month
+
+```python
+ym = YearMonth.current()
+```
+
+You can also specify a timezone using the `tz` argument:
+
+```python
+import pytz
+
+ym = YearMonth.current(tz=pytz.timezone("Asia/Tokyo"))
+```
+
+### Parsing ISO 8601 strings
 
 ```python
 ym = YearMonth.parse("2019-01")
@@ -26,16 +40,17 @@ print(ym.year) # 2019
 print(ym.month) # 1
 ```
 
-Comparisons:
+### Comparing year-months
 
 ```python
 ym1 = YearMonth(2019, 1)
 ym2 = YearMonth(2019, 2)
 
+print(ym1 == ym2) # False
 print(ym1 < ym2) # True
 ```
 
-Addition:
+### Adding or subtracting months
 
 ```python
 ym = YearMonth(2019, 1)
@@ -44,7 +59,7 @@ ym += 1
 print(ym) # 2019-02
 ```
 
-Range:
+### Iterating over a range of year-months
 
 ```python
 ym1 = YearMonth(2019, 1)
@@ -54,7 +69,7 @@ for ym in YearMonth.range(ym1, ym2):
     print(ym) # 2019-01, 2019-02, 2019-03
 ```
 
-Calculate the distance between two year-months:
+### Calculating the distance between two year-months
 
 ```python
 ym1 = YearMonth(2019, 1)

mp_yearmonth/yearmonth.py

@@ -7,7 +7,7 @@
 import re
 import calendar
 import datetime
-from typing import Iterator, Tuple
+from typing import Iterator, Optional, Tuple
 
 _FORMAT_PATTERN = re.compile(r"^(\d{4})-(\d{2})$")
 
@@ -46,7 +46,7 @@ class YearMonth:
     def __init__(self, year: int, month: int):
         """Create a new YearMonth object.
 
-        Args:
+        Arguments:
             year: The year.
             month: The month."""
         if not isinstance(year, int):
@@ -139,9 +139,13 @@ def numdays(self) -> int:
         return days
 
     @classmethod
-    def current(cls) -> "YearMonth":
-        """Return the current year and month."""
-        now = datetime.date.today()
+    def current(cls, tz: Optional[datetime.tzinfo] = None) -> "YearMonth":
+        """Return the current year and month.
+
+        Arguments:
+            tz: The timezone to use. If None, the system's local timezone is used.
+        """
+        now = datetime.datetime.now(tz)
         return cls(now.year, now.month)
 
     @classmethod

pyproject.toml

@@ -25,6 +25,7 @@ repository = "https://github.com/raymondjavaxx/mp-yearmonth"
 
 [tool.poetry.group.test.dependencies]
 pytest = "^7.4.3"
+freezegun = "^1.4.0"
 
 [build-system]
 requires = ["poetry-core"]

tests/test_yearmonth.py

@@ -1,5 +1,6 @@
 import pytest
 import datetime
+from freezegun import freeze_time
 from mp_yearmonth import YearMonth
 
 
@@ -102,8 +103,15 @@ def test_numdays():
     assert YearMonth(2020, 2).numdays == 29
 
 
+@freeze_time("2021-01-01")  # UTC
 def test_current():
-    assert YearMonth.current() == YearMonth.parse(str(YearMonth.current()))
+    assert YearMonth.current() == YearMonth(2021, 1)
+
+
+@freeze_time("2021-01-01")  # UTC
+def test_current_handles_timezone():
+    tz = datetime.timezone(datetime.timedelta(hours=-1))
+    assert YearMonth.current(tz) == YearMonth(2020, 12)
 
 
 def test_parse():