Simplify all the things

Author: maebert

src/shellcraft/core.py

@@ -4,8 +4,10 @@
 import os
 from builtins import str
 from copy import copy
+from typing import Any, Dict, List, Optional, Union
 
 import toml
+from pydantic import BaseModel, Field, root_validator
 
 from shellcraft.utils import to_float, to_list
 
@@ -119,86 +121,136 @@ def add(self, item):
         return new_item
 
 
-class BaseItem(object):
-    """Abstract class for new items."""
+class ResourceAmounts(BaseModel):
+    clay: int = 0
+    ore: int = 0
+    energy: int = 0
 
-    _pb = None
-    """This links the item to a serializable Protobuf message that is part of the GameState."""
+    def get(self, resource: str, default: int = 0) -> int:
+        return getattr(self, resource, default)
 
-    def __init__(self, name):
-        """Initialise an empty item."""
-        self.name = name
-        self.difficulty = 0
-        self.description = ""
-        self.prerequisites = {}
-        self.cost = {}
-        self.effects = {}
-        self.strings = {}
-
-    @classmethod
-    def from_dict(cls, name, data):
-        """Load an item from dict representation."""
-        item = cls(name)
-
-        item.description = data.get("description", "")
-        item.difficulty = data.get("difficulty", 0)
-
-        item.prerequisites = data.get("prerequisites", {})
-        item.prerequisites["items"] = to_list(item.prerequisites.get("items"))
-        item.prerequisites["research"] = to_list(item.prerequisites.get("research"))
-        item.prerequisites["triggers"] = to_list(item.prerequisites.get("triggers"))
-        item.cost = data.get("cost", {})
-        item.strings = data.get("strings", {})
-        item.effects = data.get("effects", {})
-        for effect in (
+    def __iter__(self):
+        yield ("clay", self.clay)
+        yield ("ore", self.ore)
+        yield ("energy", self.energy)
+
+
+class Prerequisites(ResourceAmounts):
+    items: List[str] = Field(default_factory=list)
+    research: List[str] = Field(default_factory=list)
+    triggers: List[str] = Field(default_factory=list)
+
+    @root_validator(pre=True)
+    def wrap_in_list(cls, values):
+        for field in ["items", "research", "triggers"]:
+            if field in values and not isinstance(values[field], (list, tuple)):
+                values[field] = [values[field]]
+        return values
+
+
+class Effects(ResourceAmounts):
+    enable_commands: List[str] = Field(default_factory=list)
+    enable_items: List[str] = Field(default_factory=list)
+    enable_resources: List[str] = Field(default_factory=list)
+    enable_research: List[str] = Field(default_factory=list)
+    events: List[str] = Field(default_factory=list)
+    triggers: List[str] = Field(default_factory=list)
+
+    # Resources to be granted
+    clay: int = 0
+    ore: int = 0
+    energy: int = 0
+
+    @root_validator(pre=True)
+    def wrap_in_list(cls, values):
+        for field in [
             "enable_commands",
             "enable_items",
             "enable_resources",
+            "enable_research",
             "events",
             "triggers",
-        ):
-            item.effects[effect] = to_list(item.effects.get(effect))
-        return item
+        ]:
+            if field in values and not isinstance(values[field], (list, tuple)):
+                values[field] = [values[field]]
+        return values
+
+
+class BaseItem(BaseModel):
+    """Pydantic base class for game items."""
+
+    name: str
+    difficulty: float = 0
+    description: str = ""
+    prerequisites: Prerequisites = Field(default_factory=Prerequisites)
+    effects: Effects = Field(default_factory=Effects)
+    cost: ResourceAmounts = Field(default_factory=ResourceAmounts)
+    strings: Dict[str, str] = Field(default_factory=dict)
+
+    # Use model_config to allow arbitrary types
+    model_config = {"arbitrary_types_allowed": True}
+
+    def __init__(self, name: str = "", **data):
+        """Initialise an item."""
+        if name:
+            data["name"] = name
+        super().__init__(**data)
+        # Set _pb as a private attribute after initialization
+        self._pb = None
+
+    @classmethod
+    def from_dict(cls, name, data):
+        """Load an item from dict representation."""
+        # Process prerequisites
+        data["name"] = name
+        return cls.model_validate(data)
 
     def __repr__(self):
         """String representation of the item."""
-        return self.name
+        return str(self.name)
+
+    def __str__(self):
+        """String representation of the item."""
+        return str(self.name)
 
     def __setattr__(self, key, value):
-        """Override attribute setter for items with attached message (Protobuf or Pydantic)."""
-        if key != "_pb" and self._pb is not None:
+        """Override attribute setter to sync with attached storage model (_pb)."""
+        # Let Pydantic handle its own fields first
+        super().__setattr__(key, value)
+
+        # Also sync to _pb if it exists and the key is a serializable field
+        if hasattr(self, "_pb") and self._pb is not None and key != "_pb":
             # Check if it's a protobuf object
             if hasattr(self._pb.__class__, "DESCRIPTOR"):
                 if key in self._pb.__class__.DESCRIPTOR.fields_by_name.keys():
                     setattr(self._pb, key, value)
-                else:
-                    self.__dict__[key] = value
             # Check if it's a Pydantic model
             elif hasattr(self._pb, "model_fields"):
                 if key in self._pb.model_fields:
                     setattr(self._pb, key, value)
-                else:
-                    self.__dict__[key] = value
-            else:
-                self.__dict__[key] = value
-        else:
-            self.__dict__[key] = value
-
-    def __getattr__(self, key):
-        """Override attribute getter for items with attached message (Protobuf or Pydantic)."""
-        if key != "_pb" and self._pb:
-            # Check if it's a protobuf object
-            if hasattr(self._pb.__class__, "DESCRIPTOR"):
-                if key in self._pb.__class__.DESCRIPTOR.fields_by_name.keys():
-                    return getattr(self._pb, key)
-            # Check if it's a Pydantic model
-            elif hasattr(self._pb, "model_fields"):
-                if key in self._pb.model_fields:
-                    return getattr(self._pb, key)
-        raise AttributeError(key)
 
-
-class BaseFactory(object):
+    def __getattribute__(self, key):
+        """Override attribute getter with fallback to storage model (_pb) if needed."""
+        try:
+            # Try to get from the Pydantic model first
+            return super().__getattribute__(key)
+        except AttributeError:
+            # If attribute doesn't exist on the Pydantic model, try _pb
+            if hasattr(self, "_pb") and self._pb and key != "_pb":
+                # Check if it's a protobuf object
+                if hasattr(self._pb.__class__, "DESCRIPTOR"):
+                    if key in self._pb.__class__.DESCRIPTOR.fields_by_name.keys():
+                        return getattr(self._pb, key)
+                # Check if it's a Pydantic model
+                elif hasattr(self._pb, "model_fields"):
+                    if key in self._pb.model_fields:
+                        return getattr(self._pb, key)
+            raise AttributeError(
+                f"'{self.__class__.__name__}' object has no attribute '{key}'"
+            )
+
+
+class BaseFactory:
     """Factory pattern to instantiate items."""
 
     FIXTURES = "collection.toml"
@@ -220,7 +272,7 @@ def __init__(self, game):
         }
         self.game = game
 
-    def get(self, item_name):
+    def get(self, item_name) -> Optional[BaseItem]:
         """Get an item instance by name."""
         if isinstance(item_name, BaseItem):
             return item_name
@@ -251,15 +303,19 @@ def available_items(self):
 
     def _resources_missing_to_craft(self, item_name):
         item = self.get(item_name)
+        assert item
+
         return {
             res: int(res_cost - self.game.resources.get(res))
-            for res, res_cost in item.cost.items()
+            for res, res_cost in item.cost
             if res_cost - self.game.resources.get(res) > 0
         }
 
     def can_afford(self, item_name):
         """Return true if we have enough resources to create an item."""
         item = self.get(item_name)
+        assert item
+
         for resource in RESOURCES:
             if item.cost.get(resource, 0) > self.game.resources.get(resource):
                 return False
@@ -269,33 +325,34 @@ def can_afford(self, item_name):
     def apply_effects(self, item_name):
         """Apply all effects of an item."""
         item = self.get(item_name)
+        assert item
 
         # Enable commands
-        for command in item.effects.get("enable_commands", []):
+        for command in item.effects.enable_commands:
             if command not in self.game.state.commands_enabled:
                 self.game.alert("You unlocked the `{}` command", command)
                 self.game.state.commands_enabled.append(command)
 
         # Enable resouces
-        for resources in item.effects.get("enable_resources", []):
+        for resources in item.effects.enable_resources:
             if resources not in self.game.state.resources_enabled:
                 self.game.alert("You can now mine *{}*.", resources)
                 self.game.state.resources_enabled.append(resources)
 
         # Enable items
-        for item_name in item.effects.get("enable_items", []):
+        for item_name in item.effects.enable_items:
             if item_name not in self.game.state.tools_enabled:
                 self.game.alert("You can now craft ${}$.", item_name)
                 self.game.state.tools_enabled.append(item_name)
 
         # Enable research
-        for research in item.effects.get("enable_research", []):
+        for research in item.effects.enable_research:
             if research not in self.game.state.research_enabled:
                 self.game.alert("You can now research @{}@.", research)
                 self.game.state.research_enabled.append(research)
 
         # Trigger flags
-        for trigger in item.effects.get("triggers", []):
+        for trigger in item.effects.triggers:
             if trigger not in self.game.state.triggers:
                 self.game.state.triggers.append(trigger)
 
@@ -311,7 +368,7 @@ def apply_effects(self, item_name):
 
         # Change mining difficulty
         for resource in RESOURCES:
-            change = item.effects.get(f"{resource}_mining_difficulty", None)
+            change = getattr(item.effects, f"{resource}_mining_difficulty", None)
             if change:
                 change = to_float(change)
                 self.game.mining_difficulty.multiply(resource, 1 - change)
@@ -320,7 +377,7 @@ def apply_effects(self, item_name):
                 )
 
         # Trigger events
-        self.game.events.trigger(*item.effects.get("events", []))
+        self.game.events.trigger(*item.effects.events)
 
     def is_available(self, item_name):
         """Return true if the prerequisites for an item are met."""
@@ -330,13 +387,13 @@ def is_available(self, item_name):
         for resource in RESOURCES:
             if self.game.resources.get(resource) < item.prerequisites.get(resource, 0):
                 return False
-        for required_item in item.prerequisites["items"]:
+        for required_item in item.prerequisites.items:
             if not self.game.has_item(required_item):
                 return False
-        for research in item.prerequisites["research"]:
+        for research in item.prerequisites.research:
             if research not in self.game.state.research_completed:
                 return False
-        for trigger in item.prerequisites["triggers"]:
+        for trigger in item.prerequisites.triggers:
             if trigger not in self.game.state.triggers:
                 return False
         return True

src/shellcraft/shellcraft.py

@@ -1,20 +1,19 @@
 # -*- coding: utf-8 -*-
 """Game Classes."""
 
-from shellcraft.tools import ToolFactory
-from shellcraft.research import ResearchFactory
-from shellcraft.tutorial import TutorialFactory
+import datetime
+import os
+from random import random
+
+from shellcraft.core import ItemProxy, ResourceProxy
 from shellcraft.events import EventFactory
-from shellcraft.missions import MissionFactory
-from shellcraft.fractions import FractionProxy
 from shellcraft.exceptions import BusyException
-
+from shellcraft.fractions import FractionProxy
 from shellcraft.game_state import GameState
-from shellcraft.core import ResourceProxy, ItemProxy
-
-from random import random
-import os
-import datetime
+from shellcraft.missions import MissionFactory
+from shellcraft.research import ResearchFactory
+from shellcraft.tools import ToolFactory
+from shellcraft.tutorial import TutorialFactory
 
 
 class Game(object):
@@ -95,7 +94,9 @@ def complete_missions(self):
     def craft(self, tool_name):
         """Craft a new tool by expending resources and time."""
         item = self.workshop.get(tool_name)
-        for resource, res_cost in item.cost.items():
+        assert item
+
+        for resource, res_cost in item.cost:
             self.resources.add(resource, -res_cost)
         self.tools.add(item)
         self._act("craft", tool_name, item.difficulty)
@@ -105,6 +106,8 @@ def craft(self, tool_name):
     def research(self, project_name):
         """Research a new project by expending time."""
         project = self.lab.get(project_name)
+        assert project
+
         self.state.research_completed.append(project.name)
         self._act("research", project_name, project.difficulty)
         self.alert("Researched {}.", project)
@@ -196,10 +199,11 @@ def _act(self, task, target, duration):
             duration = 0
 
         from shellcraft.game_state import Action
+
         self.state.action = Action(
             task=task,
             target=str(target),
-            completion=datetime.datetime.now() + datetime.timedelta(seconds=duration)
+            completion=datetime.datetime.now() + datetime.timedelta(seconds=duration),
         )
 
     @classmethod