faif · faif · Jan 25, 2021 · Jan 24, 2021 · Jan 24, 2021
diff --git a/patterns/creational/abstract_factory.py b/patterns/creational/abstract_factory.py
@@ -31,71 +31,85 @@
 """
 
 import random
+from typing import Type
 
 
-class PetShop:
+class Pet:
+    def __init__(self, name: str) -> None:
+        self.name = name
 
-    """A pet shop"""
+    def speak(self) -> None:
+        raise NotImplementedError
 
-    def __init__(self, animal_factory=None):
-        """pet_factory is our abstract factory.  We can set it at will."""
+    def __str__(self) -> str:
+        raise NotImplementedError
 
-        self.pet_factory = animal_factory
 
-    def show_pet(self):
-        """Creates and shows a pet using the abstract factory"""
+class Dog(Pet):
+    def speak(self) -> None:
+        print("woof")
+
+    def __str__(self) -> str:
+        return f"Dog<{self.name}>"
+
 
-        pet = self.pet_factory()
-        print(f"We have a lovely {pet}")
-        print(f"It says {pet.speak()}")
+class Cat(Pet):
+    def speak(self) -> None:
+        print("meow")
 
+    def __str__(self) -> str:
+        return f"Cat<{self.name}>"
 
-class Dog:
-    def speak(self):
-        return "woof"
 
-    def __str__(self):
-        return "Dog"
+class PetShop:
+
+    """A pet shop"""
 
+    def __init__(self, animal_factory: Type[Pet]) -> None:
+        """pet_factory is our abstract factory.  We can set it at will."""
 
-class Cat:
-    def speak(self):
-        return "meow"
+        self.pet_factory = animal_factory
+
+    def buy_pet(self, name: str) -> Pet:
+        """Creates and shows a pet using the abstract factory"""
 
-    def __str__(self):
-        return "Cat"
+        pet = self.pet_factory(name)
+        print(f"Here is your lovely {pet}")
+        return pet
 
 
 # Additional factories:
 
 # Create a random animal
-def random_animal():
+def random_animal(name: str) -> Pet:
     """Let's be dynamic!"""
-    return random.choice([Dog, Cat])()
+    return random.choice([Dog, Cat])(name)
 
 
 # Show pets with various factories
-def main():
+def main() -> None:
     """
     # A Shop that sells only cats
     >>> cat_shop = PetShop(Cat)
-    >>> cat_shop.show_pet()
-    We have a lovely Cat
-    It says meow
+    >>> pet = cat_shop.buy_pet("Lucy")
+    Here is your lovely Cat<Lucy>
+    >>> pet.speak()
+    meow
 
     # A shop that sells random animals
     >>> shop = PetShop(random_animal)
-    >>> for i in range(3):
-    ...    shop.show_pet()
+    >>> for name in ["Max", "Jack", "Buddy"]:
+    ...    pet = shop.buy_pet(name)
+    ...    pet.speak()
     ...    print("=" * 20)
-    We have a lovely Cat
-    It says meow
+    Here is your lovely Cat<Max>
+    meow
     ====================
-    We have a lovely Dog
-    It says woof
+    Here is your lovely Dog<Jack>
+    woof
     ====================
-    We have a lovely Dog
-    It says woof
+    Here is your lovely Dog<Buddy>
+    woof
     ====================
     """
 

diff --git a/tests/creational/test_abstract_factory.py b/tests/creational/test_abstract_factory.py
@@ -8,5 +8,6 @@ class TestPetShop(unittest.TestCase):
     def test_dog_pet_shop_shall_show_dog_instance(self):
         dog_pet_shop = PetShop(Dog)
         with patch.object(Dog, "speak") as mock_Dog_speak:
-            dog_pet_shop.show_pet()
+            pet = dog_pet_shop.buy_pet("")
+            pet.speak()
             self.assertEqual(mock_Dog_speak.call_count, 1)