From 25851e550c133919085b5b6da9eeeba7888c0c8d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micka=C3=ABl=20Schoentgen?= Date: Tue, 30 Jul 2024 12:26:45 +0200 Subject: [PATCH] feat: allow to use a custom date function for output file names --- CHANGELOG.md | 1 + check.sh | 2 +- docs/source/api.rst | 6 +++++- pyproject.toml | 1 + src/mss/base.py | 11 +++++++++-- src/tests/test_save.py | 30 +++++++++++++++++++++++++++++- 6 files changed, 46 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e0e18ea6..625b4c11 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ See Git checking messages for full history. ## 9.0.2 (2023/xx/xx) - 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/check.sh b/check.sh index 7bb90ae2..8028ab20 100755 --- a/check.sh +++ b/check.sh @@ -4,8 +4,8 @@ # set -eu -python -m ruff --fix docs src python -m ruff format docs src +python -m ruff check --fix docs src # "--platform win32" to not fail on ctypes.windll (it does not affect the overall check on other OSes) python -m mypy --platform win32 src docs/source/examples 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 5a0d1252..30d8b3f7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,6 +70,7 @@ mss = "mss.__main__:main" [project.optional-dependencies] test = [ + "freezegun", "numpy", "pillow", "pytest", diff --git a/src/mss/base.py b/src/mss/base.py index 4495bf14..d09a5ba7 100644 --- a/src/mss/base.py +++ b/src/mss/base.py @@ -28,6 +28,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.""" @@ -130,6 +134,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. @@ -153,6 +158,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). """ @@ -164,7 +171,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) @@ -180,7 +187,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 a46a4fb2..34cb2b81 100644 --- a/src/tests/test_save.py +++ b/src/tests/test_save.py @@ -2,9 +2,10 @@ Source: https://github.com/BoboTiG/python-mss. """ import os.path -from datetime import datetime +from datetime import datetime, timedelta import pytest +from freezegun import freeze_time from mss import mss try: @@ -16,6 +17,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)) @@ -78,3 +82,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)