diff --git a/fimfarchive/stampers.py b/fimfarchive/stampers.py index 00a1e23..877d1e9 100644 --- a/fimfarchive/stampers.py +++ b/fimfarchive/stampers.py @@ -22,12 +22,13 @@ Stampers for Fimfarchive. # -from typing import Any, Dict, Optional, Set, Type, TypeVar +from typing import Any, Dict, Set import arrow -from fimfarchive.flavors import Flavor, UpdateStatus +from fimfarchive.flavors import UpdateStatus from fimfarchive.stories import Story +from fimfarchive.utils import find_flavor __all__ = ( @@ -36,9 +37,6 @@ __all__ = ( ) -F = TypeVar('F', bound=Flavor) - - class Stamper: """ Adds archive-related information to stories. @@ -90,23 +88,6 @@ class UpdateStamper(Stamper): }, } - def find_flavor(self, story: Story, flavor: Type[F]) -> Optional[F]: - """ - Searches for a flavor of a specific type. - - Args: - story: The story to search in. - flavor: The flavor type to find. - - Returns: - A flavor of the desired type, or None. - """ - for current in story.flavors: - if isinstance(current, flavor): - return current - - return None - def __call__(self, story: Story) -> None: """ Applies modification dates to a story. @@ -115,7 +96,7 @@ class UpdateStamper(Stamper): story: The story to stamp. """ timestamp = arrow.utcnow().isoformat() - flavor = self.find_flavor(story, UpdateStatus) + flavor = find_flavor(story, UpdateStatus) archive = self.get_archive(story) archive['date_checked'] = timestamp diff --git a/fimfarchive/utils.py b/fimfarchive/utils.py index 77c0a6f..e538cc9 100644 --- a/fimfarchive/utils.py +++ b/fimfarchive/utils.py @@ -25,15 +25,22 @@ Various utilities. import json import os import shutil -from typing import Dict, Any +from typing import Any, Dict, Optional, Type, TypeVar + +from fimfarchive.flavors import Flavor +from fimfarchive.stories import Story __all__ = ( 'Empty', 'PersistedDict', + 'find_flavor', ) +F = TypeVar('F', bound=Flavor) + + class EmptyMeta(type): """ Meta-class for Empty. @@ -101,3 +108,21 @@ class PersistedDict(Dict[str, Any]): if os.path.exists(self.temp): os.remove(self.temp) + + +def find_flavor(story: Story, flavor: Type[F]) -> Optional[F]: + """ + Searches for a flavor of a specific type. + + Args: + story: The story to search in. + flavor: The flavor type to find. + + Returns: + A flavor of the desired type, or None. + """ + for current in story.flavors: + if isinstance(current, flavor): + return current + + return None diff --git a/tests/test_stampers.py b/tests/test_stampers.py index e9cf058..0891491 100644 --- a/tests/test_stampers.py +++ b/tests/test_stampers.py @@ -63,8 +63,8 @@ class TestStamper: meta['archive'] = original archive = stamper.get_archive(story) - assert archive is original assert meta['archive'] is original + assert archive is original class TestUpdateStamper: diff --git a/tests/test_utils.py b/tests/test_utils.py index 58bf327..c03baa7 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -27,7 +27,8 @@ import os import pytest -from fimfarchive.utils import Empty, PersistedDict +from fimfarchive.flavors import DataFormat, MetaFormat, MetaPurity +from fimfarchive.utils import find_flavor, Empty, PersistedDict class TestEmpty: @@ -189,3 +190,33 @@ class TestPersistedDict: data = PersistedDict(tmpfile, default=extra) assert dict(data) == sample + + +class TestFindFlavor: + """ + find_flavor tests. + """ + + @pytest.fixture + def story(self, story): + """ + Returns a meta-flavored story. + """ + return story.merge(flavors=[ + MetaFormat.BETA, + MetaPurity.CLEAN, + ]) + + def test_present_flavor(self, story): + """ + Tests flavor is returned when present. + """ + found = find_flavor(story, MetaFormat) + assert found == MetaFormat.BETA + + def test_missing_flavor(self, story): + """ + Tests None is returned when flavor is missing. + """ + found = find_flavor(story, DataFormat) + assert found is None