ooad1-6 Exercise 2: Replace Type Tests With Polymorphism ============================================================= shipping.py (the starting point, two if-chains): """Exercise 2, before: shipping cost with a chain of type tests. Behaviour to keep: standard 400p + 50p per kg, express 900p + 120p per kg, pickup 0.""" def shipping_cost(method: str, kilos: float) -> int: if method == "standard": return 400 + int(50 * kilos) elif method == "express": return 900 + int(120 * kilos) elif method == "pickup": return 0 else: raise ValueError(f"unknown shipping method {method!r}") def describe(method: str) -> str: if method == "standard": return "Standard delivery, 3 to 5 days" elif method == "express": return "Express delivery, next day" elif method == "pickup": return "Collect from the shop" else: raise ValueError(f"unknown shipping method {method!r}") shipping_poly.py (the replacement): """Exercise 2, after: one class per shipping method (Polymorphism). Same behaviour, no type tests.""" from abc import ABC, abstractmethod class Shipping(ABC): @abstractmethod def cost(self, kilos: float) -> int: ... @abstractmethod def description(self) -> str: ... class Standard(Shipping): def cost(self, kilos: float) -> int: return 400 + int(50 * kilos) def description(self) -> str: return "Standard delivery, 3 to 5 days" class Express(Shipping): def cost(self, kilos: float) -> int: return 900 + int(120 * kilos) def description(self) -> str: return "Express delivery, next day" class Pickup(Shipping): def cost(self, kilos: float) -> int: return 0 def description(self) -> str: return "Collect from the shop" METHODS = {"standard": Standard(), "express": Express(), "pickup": Pickup()} def shipping_cost(method: str, kilos: float) -> int: return lookup(method).cost(kilos) def describe(method: str) -> str: return lookup(method).description() def lookup(method: str) -> Shipping: if method not in METHODS: raise ValueError(f"unknown shipping method {method!r}") return METHODS[method] ex2.py (the check): import ast import pathlib import unittest import shipping as before import shipping_poly as after CASES = [("standard", 0), ("standard", 2.5), ("express", 0), ("express", 3.2), ("pickup", 10)] class Same(unittest.TestCase): def test_costs_and_descriptions_match(self): for method, kilos in CASES: self.assertEqual(before.shipping_cost(method, kilos), after.shipping_cost(method, kilos)) self.assertEqual(before.describe(method), after.describe(method)) def test_unknown_method_refused_by_both(self): for module in (before, after): with self.assertRaises(ValueError): module.shipping_cost("drone", 1) with self.assertRaises(ValueError): module.describe("drone") def branches(path: str) -> int: tree = ast.parse(pathlib.Path(path).read_text(encoding="utf8")) return sum(isinstance(n, ast.If) for n in ast.walk(tree)) if __name__ == "__main__": result = unittest.TextTestRunner(verbosity=0).run(unittest.defaultTestLoader.loadTestsFromTestCase(Same)) print("same behaviour:", result.wasSuccessful(), f"({result.testsRun} tests)") print("if/elif branches before:", branches("shipping.py"), "| after:", branches("shipping_poly.py")) # adding a 'drone' method: # before: edit shipping_cost AND describe (two places, each must remember) # after : write one class and add one line to METHODS Run with: py ex2.py Output: same behaviour: True (2 tests) if/elif branches before: 6 | after: 1 COUNTING THE COST OF A 'drone' METHOD before: edit shipping_cost() and describe(): two existing functions, and if you edit one and forget the other, 'drone' is priced but cannot be described (or the reverse), and nothing tells you until someone calls the forgotten one. after : write one class (Drone) that must define cost() and description() (Python refuses to create it otherwise), and add one line to METHODS. Existing code is untouched. The 'after' still contains one if: the check that the method name is known. That is validation at the edge, not a type test on a tag, and it is the only branch left. WHY THIS WORKS AS AN ANSWER --------------------------- The behaviour check compares the old and new code on five cases and on an unknown method, so the refactoring is safe. The branch count (6 down to 1) measures the type tests removed. The real benefit is the abstract base class: it turns 'remember to update both places' into 'the program will not start until you have'.