From 9ca187234355045e8e3d831e0ee56a7ef42fd7e1 Mon Sep 17 00:00:00 2001
From: Joakim Soderlund <joakim.soderlund@gmail.com>
Date: Sun, 16 Jun 2019 13:52:19 +0200
Subject: [PATCH] Add flavor stamper

---
 fimfarchive/stampers.py | 25 +++++++++++++++++++++++--
 tests/test_stampers.py  | 35 +++++++++++++++++++++++++++++++++--
 2 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/fimfarchive/stampers.py b/fimfarchive/stampers.py
index f5ead18..c655297 100644
--- a/fimfarchive/stampers.py
+++ b/fimfarchive/stampers.py
@@ -26,7 +26,7 @@ from typing import Any, Callable, Dict, Optional, Set
 
 import arrow
 
-from fimfarchive.flavors import UpdateStatus
+from fimfarchive.flavors import Flavor, UpdateStatus
 from fimfarchive.stories import Story
 from fimfarchive.utils import find_flavor
 
@@ -108,6 +108,27 @@ class UpdateStamper(Stamper):
                 archive[key] = None
 
 
+class FlavorStamper(Stamper):
+    """
+    Adds flavors to stories.
+    """
+
+    def __init__(self, mapper: Callable[[Story], Optional[Flavor]]) -> None:
+        """
+        Constructor.
+
+        Args:
+            mapper: Callable returning the flavor to stamp with.
+        """
+        self.map = mapper
+
+    def __call__(self, story: Story) -> None:
+        flavor = self.map(story)
+
+        if flavor:
+            story.flavors.add(flavor)
+
+
 class PathStamper(Stamper):
     """
     Adds archive paths to stories.
@@ -118,7 +139,7 @@ class PathStamper(Stamper):
         Constructor.
 
         Args:
-            mapper: Callable returning the path to stamp.
+            mapper: Callable returning the path to stamp with.
         """
         self.map = mapper
 
diff --git a/tests/test_stampers.py b/tests/test_stampers.py
index 3a77df9..64a874d 100644
--- a/tests/test_stampers.py
+++ b/tests/test_stampers.py
@@ -28,9 +28,11 @@ from typing import Dict
 import arrow
 import pytest
 
-from fimfarchive.flavors import UpdateStatus
+from fimfarchive.flavors import DataFormat, MetaFormat, UpdateStatus
 from fimfarchive.mappers import StaticMapper
-from fimfarchive.stampers import Stamper, PathStamper, UpdateStamper
+from fimfarchive.stampers import (
+    Stamper, FlavorStamper, PathStamper, UpdateStamper,
+)
 
 
 class TestStamper:
@@ -190,6 +192,35 @@ class TestUpdateStamper:
         assert archive['date_updated'] == prev
 
 
+class TestFlavorStamper:
+    """
+    FlavorStamper tests.
+    """
+
+    @pytest.mark.parametrize('value', (None, ''))
+    def test_ignored_blank_value(self, story, value):
+        """
+        Tests blank values are ignored.
+        """
+        stamp = FlavorStamper(StaticMapper(value))
+        stamp(story)
+
+        assert value not in story.flavors
+
+    @pytest.mark.parametrize('value', (
+        DataFormat.EPUB,
+        MetaFormat.ALPHA,
+    ))
+    def test_stamped_value(self, story, value):
+        """
+        Tests values are stamped to stories.
+        """
+        stamp = FlavorStamper(StaticMapper(value))
+        stamp(story)
+
+        assert value in story.flavors
+
+
 class TestPathStamper:
     """
     PathStamper tests.