Add inventory arithmetic

elite_engineering/materials.py

@@ -177,7 +177,15 @@ class Inventory(collections.UserDict):
     def __setitem__(self, key, value) -> None:
         if key not in self.data:
             raise KeyError(key)
+        if value < 0:
+            raise ValueError(value)
         super().__setitem__(key, value)
 
-    def has(self, items: typing.Dict[str, int]) -> bool:
+    def has(self, items: typing.Union[typing.Dict[str, int], "Inventory"]) -> bool:
         return all([self[name] >= quantity for name, quantity in items.items()])
+
+    def __add__(self, other: "Inventory") -> "Inventory":
+        return Inventory({name: self[name] + qty for name, qty in other.items()})
+
+    def __sub__(self, other: "Inventory") -> "Inventory":
+        return Inventory({name: self[name] - qty for name, qty in other.items()})

tests/test_materials.py

@@ -109,6 +109,11 @@ class InventoryTests(unittest.TestCase):
             materials.Inventory({"Carbon": 5, "Pizza": 12})
         self.assertEqual(exc.exception.args, ("Pizza",))
 
+    def test_create_inventory_with_invalid_quantities(self) -> None:
+        with self.assertRaises(ValueError) as exc:
+            materials.Inventory({"Carbon": -5})
+        self.assertEqual(exc.exception.args, (-5,))
+
     def test_inventory_has_materials(self) -> None:
         inventory = materials.Inventory({"Carbon": 5, "Shield Emitters": 3})
         self.assertTrue(inventory.has({"Carbon": 1, "Shield Emitters": 3}))
@@ -123,3 +128,22 @@ class InventoryTests(unittest.TestCase):
         with self.assertRaises(KeyError) as exc:
             inventory.has({"Carbon": 5, "Pizza": 12})
         self.assertEqual(exc.exception.args, ("Pizza",))
+
+    def test_inventory_addition(self) -> None:
+        a = materials.Inventory({"Carbon": 5, "Shield Emitters": 3})
+        b = materials.Inventory({"Iron": 2, "Shield Emitters": 3})
+        expected = materials.Inventory({"Carbon": 5, "Iron": 2, "Shield Emitters": 6})
+        self.assertEqual(expected, a + b)
+
+    def test_inventory_subtraction(self) -> None:
+        a = materials.Inventory({"Carbon": 5, "Iron": 2, "Shield Emitters": 6})
+        b = materials.Inventory({"Carbon": 5, "Shield Emitters": 3})
+        expected = materials.Inventory({"Iron": 2, "Shield Emitters": 3})
+        self.assertEqual(expected, a - b)
+
+    def test_inventory_subtraction_cannot_result_in_negative_quantities(self) -> None:
+        a = materials.Inventory({"Carbon": 5, "Iron": 2, "Shield Emitters": 6})
+        b = materials.Inventory({"Carbon": 6, "Shield Emitters": 3})
+        with self.assertRaises(ValueError) as exc:
+            a - b
+        self.assertEqual(exc.exception.args, (-1,))