diff --git a/CHANGELOG.md b/CHANGELOG.md index ddb1a8fa..f40df2eb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ See Git checking messages for full history. - added support for Python 3.13 - leveled up the packaging using `hatchling` - used `ruff` to lint the code base (#275) +- MSS: allow to use a custom date function for output file names (#276) - MSS: minor optimization when using an output file format without date (#275) - MSS: fixed `Pixel` model type (#274) - CI: automated release publishing on tag creation diff --git a/docs/source/api.rst b/docs/source/api.rst index cef76dab..d297968a 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -168,12 +168,16 @@ Methods *monitor* can be a ``tuple`` like ``PIL.Image.grab()`` accepts, it will be converted to the appropriate ``dict``. - .. method:: save([mon=1], [output='mon-{mon}.png'], [callback=None]) + .. method:: save([mon=1], [output='mon-{mon}.png'], [callback=None], [date_fn=None]) :param int mon: the monitor's number. :param str output: the output's file name. :type callback: callable or None :param callback: callback called before saving the screen shot to a file. Takes the *output* argument as parameter. + :type callback: callable or None + :param callback: callback called before saving the screen shot to a file. Takes the *output* argument as parameter. + Function returning a `datetime` object, + used to format the date in output file names. :rtype: iterable :return: Created file(s). diff --git a/pyproject.toml b/pyproject.toml index c8065c39..b3cce0f3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,6 +71,7 @@ mss = "mss.__main__:main" [project.optional-dependencies] test = [ + "freezegun", "numpy ; sys_platform == 'windows' and python_version >= '3.13'", "pillow", "pytest", diff --git a/src/mss/base.py b/src/mss/base.py index 4438693f..a2cb5236 100644 --- a/src/mss/base.py +++ b/src/mss/base.py @@ -29,6 +29,10 @@ lock = Lock() +def date_function() -> datetime: + return datetime.now(UTC) + + class MSSBase(metaclass=ABCMeta): """This class will be overloaded by a system specific one.""" @@ -131,6 +135,7 @@ def save( mon: int = 0, output: str = "monitor-{mon}.png", callback: Callable[[str], None] | None = None, + date_fn: Callable[[], datetime] = date_function, ) -> Iterator[str]: """Grab a screen shot and save it to a file. @@ -154,6 +159,8 @@ def save( :param callable callback: Callback called before saving the screen shot to a file. Take the `output` argument as parameter. + :param callable date_fn: Function returning a `datetime` object, + used to format the date in output file names. :return generator: Created file(s). """ @@ -165,7 +172,7 @@ def save( if mon == 0: # One screen shot by monitor for idx, monitor in enumerate(monitors[1:], 1): - fname = output.format(mon=idx, date=datetime.now(UTC) if "{date" in output else None, **monitor) + fname = output.format(mon=idx, date=date_fn() if "{date" in output else None, **monitor) if callable(callback): callback(fname) sct = self.grab(monitor) @@ -181,7 +188,7 @@ def save( msg = f"Monitor {mon!r} does not exist." raise ScreenShotError(msg) from exc - output = output.format(mon=mon, date=datetime.now(UTC) if "{date" in output else None, **monitor) + output = output.format(mon=mon, date=date_fn() if "{date" in output else None, **monitor) if callable(callback): callback(output) sct = self.grab(monitor) diff --git a/src/tests/test_save.py b/src/tests/test_save.py index 3ae0f6f8..2f50e80a 100644 --- a/src/tests/test_save.py +++ b/src/tests/test_save.py @@ -3,9 +3,10 @@ """ import os.path -from datetime import datetime +from datetime import datetime, timedelta import pytest +from freezegun import freeze_time from mss import mss try: @@ -17,6 +18,9 @@ UTC = timezone.utc +FROZEN_TIME = "2024-02-21 16:35:20" + + def test_at_least_2_monitors() -> None: with mss(display=os.getenv("DISPLAY")) as sct: assert list(sct.save(mon=0)) @@ -79,3 +83,27 @@ def test_output_format_date_custom() -> None: filename = sct.shot(mon=1, output=fmt) assert filename == fmt.format(date=datetime.now(tz=UTC)) assert os.path.isfile(filename) + + +@freeze_time(FROZEN_TIME) +def test_output_format_custom_date_function() -> None: + def custom_date() -> datetime: + return datetime.now(tz=UTC) + timedelta(days=6) + + fmt = "{date}.png" + with mss(display=os.getenv("DISPLAY")) as sct: + filename = sct.shot(mon=1, output=fmt, date_fn=custom_date) + assert filename == "2024-02-27 16:35:20+00:00.png" + assert os.path.isfile(filename) + + +@freeze_time(FROZEN_TIME) +def test_output_format_date_custom_and_custom_date_function() -> None: + def custom_date() -> datetime: + return datetime.now(tz=UTC) + timedelta(days=6) + + fmt = "{date:%Y-%m-%d}.png" + with mss(display=os.getenv("DISPLAY")) as sct: + filename = sct.shot(mon=1, output=fmt, date_fn=custom_date) + assert filename == "2024-02-27.png" + assert os.path.isfile(filename)