From 64317a7019c63cb136c0e49374f0131a9ac8c24f Mon Sep 17 00:00:00 2001 From: Correl Roush Date: Sat, 22 Jan 2022 02:35:27 -0500 Subject: [PATCH] Add a Materials collection --- elite_engineering/materials.py | 41 +++++++++++++++++++- tests/test_materials.py | 68 +++++++++++++++++++++------------- 2 files changed, 81 insertions(+), 28 deletions(-) diff --git a/elite_engineering/materials.py b/elite_engineering/materials.py index 5c50648..1ac5660 100644 --- a/elite_engineering/materials.py +++ b/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}) diff --git a/tests/test_materials.py b/tests/test_materials.py index 67ddeb0..799437a 100644 --- a/tests/test_materials.py +++ b/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()))