From 9ef9239ab6ab6b20b07f8f5a288707e4cd12b4f4 Mon Sep 17 00:00:00 2001 From: Correl Roush Date: Sat, 22 Jan 2022 01:55:08 -0500 Subject: [PATCH] Add inventory arithmetic --- elite_engineering/materials.py | 10 +++++++++- tests/test_materials.py | 24 ++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/elite_engineering/materials.py b/elite_engineering/materials.py index 93ddb9e..5c50648 100644 --- a/elite_engineering/materials.py +++ b/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()}) diff --git a/tests/test_materials.py b/tests/test_materials.py index c2deb9a..67ddeb0 100644 --- a/tests/test_materials.py +++ b/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,))