Add a Materials collection

elite_engineering/materials.py

@@ -169,13 +169,18 @@ materials: typing.Dict[str, Material] = {
 
 
 class Inventory(collections.UserDict):
-    def __init__(self, items: typing.Optional[typing.Dict[str, int]] = None) -> None:
+    """A mapping of every available material with an associated quantity."""
+
+    def __init__(
+        self,
+        items: typing.Optional[typing.Union[typing.Dict[str, int], "Inventory"]] = None,
+    ) -> None:
         self.data = {name: 0 for name in materials}
         if items:
             self.update(items)
 
     def __setitem__(self, key, value) -> None:
-        if key not in self.data:
+        if key not in materials:
             raise KeyError(key)
         if value < 0:
             raise ValueError(value)
@@ -189,3 +194,35 @@ class Inventory(collections.UserDict):
 
     def __sub__(self, other: "Inventory") -> "Inventory":
         return Inventory({name: self[name] - qty for name, qty in other.items()})
+
+
+class Materials(Inventory):
+    """Same as Inventory, but omits zeroed materials."""
+
+    def __init__(
+        self,
+        items: typing.Optional[typing.Union[typing.Dict[str, int], "Inventory"]] = None,
+    ) -> None:
+        self.data: typing.Dict[str, int] = {}
+        if items:
+            self.update(items)
+
+    def __setitem__(self, key: str, value: int) -> None:
+        if key not in materials:
+            raise KeyError(key)
+        if value < 0:
+            raise ValueError(value)
+        if value == 0:
+            return
+        super().__setitem__(key, value)
+
+    def __getitem__(self, key: str) -> int:
+        if key not in materials:
+            raise KeyError(key)
+        return self.data.get(key, 0)
+
+    def __add__(self, other: "Inventory") -> "Materials":
+        return Materials({name: self[name] + other[name] for name in materials})
+
+    def __sub__(self, other: "Inventory") -> "Materials":
+        return Materials({name: self[name] - other[name] for name in materials})

tests/test_materials.py

@@ -94,56 +94,72 @@ class TradeCalculationTests(unittest.TestCase):
 
 
 class InventoryTests(unittest.TestCase):
-    def test_create_empty_inventory(self) -> None:
-        inventory = materials.Inventory()
-        self.assertEqual(set(inventory.keys()), set(materials.materials.keys()))
+    def make(
+        self, items: typing.Optional[typing.Dict[str, int]] = None
+    ) -> "materials.Inventory":
+        return materials.Inventory(items)
 
-    def test_create_populated_inventory(self) -> None:
-        inventory = materials.Inventory({"Carbon": 5, "Shield Emitters": 3})
+    def test_create_empty(self) -> None:
+        inventory = self.make()
+        self.assertEqual(set(materials.materials.keys()), set(inventory.keys()))
+
+    def test_create_populated(self) -> None:
+        inventory = self.make({"Carbon": 5, "Shield Emitters": 3})
         self.assertEqual(5, inventory["Carbon"])
         self.assertEqual(3, inventory["Shield Emitters"])
         self.assertEqual(8, sum(inventory.values()))
 
-    def test_create_inventory_with_invalid_materials(self) -> None:
+    def test_create_with_invalid_materials(self) -> None:
         with self.assertRaises(KeyError) as exc:
-            materials.Inventory({"Carbon": 5, "Pizza": 12})
+            self.make({"Carbon": 5, "Pizza": 12})
         self.assertEqual(exc.exception.args, ("Pizza",))
 
-    def test_create_inventory_with_invalid_quantities(self) -> None:
+    def test_create_with_invalid_quantities(self) -> None:
         with self.assertRaises(ValueError) as exc:
-            materials.Inventory({"Carbon": -5})
+            self.make({"Carbon": -5})
         self.assertEqual(exc.exception.args, (-5,))
 
-    def test_inventory_has_materials(self) -> None:
-        inventory = materials.Inventory({"Carbon": 5, "Shield Emitters": 3})
+    def test_has_materials(self) -> None:
+        inventory = self.make({"Carbon": 5, "Shield Emitters": 3})
         self.assertTrue(inventory.has({"Carbon": 1, "Shield Emitters": 3}))
 
-    def test_inventory_lacks_materials(self) -> None:
-        inventory = materials.Inventory({"Carbon": 5, "Shield Emitters": 3})
+    def test_lacks_materials(self) -> None:
+        inventory = self.make({"Carbon": 5, "Shield Emitters": 3})
         self.assertTrue(not inventory.has({"Nickel": 1}))
         self.assertTrue(not inventory.has({"Carbon": 6, "Shield Emitters": 3}))
 
-    def test_inventory_has_raises_on_invalid_materials(self) -> None:
-        inventory = materials.Inventory()
+    def test_has_raises_on_invalid_materials(self) -> None:
+        inventory = self.make()
         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})
+    def test_addition(self) -> None:
+        a = self.make({"Carbon": 5, "Shield Emitters": 3})
+        b = self.make({"Iron": 2, "Shield Emitters": 3})
+        expected = self.make({"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})
+    def test_subtraction(self) -> None:
+        a = self.make({"Carbon": 5, "Iron": 2, "Shield Emitters": 6})
+        b = self.make({"Carbon": 5, "Shield Emitters": 3})
+        expected = self.make({"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})
+    def test_subtraction_cannot_result_in_negative_quantities(self) -> None:
+        a = self.make({"Carbon": 5, "Iron": 2, "Shield Emitters": 6})
+        b = self.make({"Carbon": 6, "Shield Emitters": 3})
         with self.assertRaises(ValueError) as exc:
             a - b
         self.assertEqual(exc.exception.args, (-1,))
+
+
+class MaterialsTests(InventoryTests):
+    def make(
+        self, items: typing.Optional[typing.Dict[str, int]] = None
+    ) -> "materials.Materials":
+        return materials.Materials(items)
+
+    def test_create_empty(self) -> None:
+        empty = self.make()
+        self.assertEqual(set(), set(empty.keys()))