feat(Artifact): add ArtifactSet type

This is an Enum of ArtifactSetData instances containing data for the set_num (set bonuses), rarities, and valid slots for the set (this is for special cases like the Prayers sets that are only circlets).

src/artipy/artifacts/artifact.py

@@ -4,12 +4,10 @@
 
 from artipy import UPGRADE_STEP
 from artipy.stats import MainStat, SubStat
-from artipy.types import ArtifactSlot, StatType
+from artipy.types import ArtifactSet, ArtifactSlot
 
 from .upgrade_strategy import AddStatStrategy, UpgradeStatStrategy, UpgradeStrategy
 
-PLACEHOLDER_MAINSTAT = MainStat(StatType.HP, 0)
-
 
 class Artifact:
     """Class representing an artifact in Genshin Impact."""
@@ -21,19 +19,17 @@ def __init__(self) -> None:
         self._substats: list[SubStat] = []
         self._level: int = 0
         self._rarity: int = 0
-        self._set: str = ""
-        self._slot: ArtifactSlot | str = ""
+        self._set: Optional[ArtifactSet] = None
+        self._slot: Optional[ArtifactSlot] = None
 
     @property
-    def mainstat(self) -> MainStat:
+    def mainstat(self) -> Optional[MainStat]:
         """The mainstat of the artifact. Return a placeholder mainstat if the mainstat
         is None.
 
         Returns:
-            MainStat: The mainstat of the artifact.
+            Optional[MainStat]: The mainstat of the artifact.
         """
-        if self._mainstat is None:
-            return PLACEHOLDER_MAINSTAT
         return self._mainstat
 
     @mainstat.setter
@@ -114,46 +110,47 @@ def rarity(self, rarity: int) -> None:
             rarity (int): The rarity to set.
         """
         self._rarity = rarity
-        stats: list[MainStat | SubStat] = [self.mainstat, *self.substats]
+        stats: list[Optional[MainStat] | SubStat] = [self.mainstat, *self.substats]
         for stat in stats:
             try:
-                stat.rarity = rarity
+                if stat is not None:
+                    stat.rarity = rarity
             except AttributeError:
                 continue
 
     @property
-    def artifact_set(self) -> str:
+    def artifact_set(self) -> Optional[ArtifactSet]:
         """The artifact set of the artifact.
 
         Returns:
-            str: The artifact set of the artifact.
+            Optional[ArtifactSet]: The artifact set of the artifact.
         """
         return self._set
 
     @artifact_set.setter
-    def artifact_set(self, artifact_set: str) -> None:
+    def artifact_set(self, artifact_set: ArtifactSet) -> None:
         """Set the artifact set of the artifact.
 
         Args:
-            artifact_set (str): The artifact set to set.
+            artifact_set (ArtifactSet): The artifact set to set.
         """
         self._set = artifact_set
 
     @property
-    def artifact_slot(self) -> str | ArtifactSlot:
+    def artifact_slot(self) -> Optional[ArtifactSlot]:
         """The artifact slot of the artifact.
 
         Returns:
-            str | ArtifactSlot: The artifact slot of the artifact.
+            Optional[ArtifactSlot]: The artifact slot of the artifact.
         """
         return self._slot
 
     @artifact_slot.setter
-    def artifact_slot(self, slot: str | ArtifactSlot) -> None:
+    def artifact_slot(self, slot: ArtifactSlot) -> None:
         """Set the artifact slot of the artifact.
 
         Args:
-            slot (str | ArtifactSlot): The artifact slot to set.
+            slot (ArtifactSlot): The artifact slot to set.
         """
         self._slot = slot
 

src/artipy/artifacts/builder.py

@@ -6,7 +6,7 @@
 
 from artipy import MAX_RARITY, UPGRADE_STEP
 from artipy.stats import MainStat, SubStat
-from artipy.types import ArtifactSlot, StatType
+from artipy.types import ArtifactSet, ArtifactSlot, StatType
 
 from .artifact import Artifact
 from .upgrade_strategy import AddStatStrategy
@@ -15,15 +15,18 @@
 class ArtifactBuilder:
     """Builder class for creating an Artifact object.
 
-    Options:
-    - with_mainstat: Set the mainstat of the artifact
-    - with_substat: Add a substat to the artifact
-    - with_substats: Add multiple substats to the artifact
-    - with_level: Set the level of the artifact
-    - with_rarity: Set the rarity of the artifact
-    - with_set: Set the artifact set
-    - with_slot: Set the artifact slot
-    - build: Build the artifact object based on the parameters passed into the builder
+    Parameters:
+        - with_mainstat: Set the mainstat of the artifact
+        - with_substat: Add a substat to the artifact
+        - with_substats: Add multiple substats to the artifact
+        - with_level: Set the level of the artifact
+        - with_rarity: Set the rarity of the artifact
+        - with_set: Set the artifact set
+        - with_slot: Set the artifact slot
+
+    Methods:
+        - build: Build the artifact object based on the parameters passed into the
+                 builder.
     """
 
     def __init__(self) -> None:
@@ -32,33 +35,33 @@ def __init__(self) -> None:
     def with_mainstat(
         self, stat: StatType, value: float | int = 0
     ) -> "ArtifactBuilder":
-        """Set the mainstat of the artifact. The value will be set based on the level of
-        the artifact.
-
-        :param stat: The mainstat to set.
-        :type stat: StatType
-        :param value: The value to set, defaults to 0
-        :type value: float | int, optional
-        :return: The artifact builder object
-        :rtype: ArtifactBuilder
+        """Set the mainstat of the artifact.
+
+        Args:
+            stat (StatType): The mainstat to set.
+            value (float | int, optional): The value of the mainstat. Defaults to 0.
+
+        Returns:
+            ArtifactBuilder: The artifact builder object
         """
         self._artifact.mainstat = MainStat(stat, value)
         self._artifact.mainstat.set_value_by_level(self._artifact.level)
         return self
 
     def with_substat(self, stat: StatType, value: float | int) -> "ArtifactBuilder":
-        """Add a substat to the artifact.
-
-        :param stat: The substat to add.
-        :type stat: StatType
-        :param value: The value of the substat.
-        :type value: float | int
-        :raises ValueError: If the number of substats exceeds the rarity of the artifact
-        :return: The artifact builder object
-        :rtype: ArtifactBuilder
+        """Set a single substat of the artifact.
+
+        Args:
+            stat (StatType): The substat to set.
+            value (float | int): The value of the substat.
+
+        Raises:
+            ValueError: If the substats are already full
+
+        Returns:
+            ArtifactBuilder: The artifact builder object
         """
         if (rarity := self._artifact.rarity) > 0:
-            # Constraint: The number of substats can't exceed the rarity of the artifact
             if len(self._artifact.substats) >= rarity - 1:
                 raise ValueError("Substats are already full.")
         self._artifact.add_substat(SubStat(stat, value))
@@ -70,22 +73,24 @@ def with_substats(
         *,
         amount: int = 0,
     ) -> "ArtifactBuilder":
-        """Add multiple substats to the artifact. If substats is not provided and an
-        amount is provided, random substats will be generated.
-
-        :param substats: The substats to add, defaults to None
-        :type substats: list[tuple[StatType, float | int]], optional
-        :param amount: The amount of substats to add (if none are provided),
-                       defaults to 0
-        :type amount: int, optional
-        :raises ValueError: If the amount is not within the valid range
-        :return: The artifact builder object
-        :rtype: ArtifactBuilder
+        """Set the substats of the artifact. If no substats are provided, generate
+        random substats based on the rarity of the artifact.
+
+        Args:
+            substats (Optional[list[tuple[StatType, float  |  int]]], optional): The
+                substats to set. Defaults to None.
+            amount (int, optional): The amount of stats to generate. Defaults to 0.
+
+        Raises:
+            ValueError: If the amount is not within the valid range
+            ValueError: If the number of substats exceeds the rarity of the artifact
+
+        Returns:
+            ArtifactBuilder: The artifact builder object
         """
         rarity = self._artifact.rarity
         if substats is None:
             substats = []
-            # Constraint: The number of substats cannot exceed artifact rarity - 1
             valid_range = range(1, rarity)
             if amount not in valid_range:
                 raise ValueError(
@@ -96,60 +101,61 @@ def with_substats(
                 new_stat = AddStatStrategy().pick_stat(self._artifact)
                 self._artifact.add_substat(new_stat)
         elif len(substats) > rarity - 1 and rarity > 0:
-            # Constraint: The number of substats cannot exceed artifact rarity
             raise ValueError("Too many substats provided.")
         else:
             self._artifact.substats = [SubStat(*spec) for spec in substats]
 
         return self
 
     def with_level(self, level: int) -> "ArtifactBuilder":
-        """Set the level of the artifact. Set any values that are dependent on the
-        level.
-
-        :param level: The level to set.
-        :type level: int
-        :raises ValueError: If the level is not within the valid range
-        :raises ValueError: If the number of substats exceeds the rarity of the artifact
-        :return: The artifact builder object
-        :rtype: ArtifactBuilder
+        """Set the level of the artifact. The level determines the value of the
+        mainstat.
+
+        Args:
+            level (int): The level to set.
+
+        Raises:
+            ValueError: If the level is not within the valid range
+            ValueError: If there is a substat length mismatch with the rarity
+
+        Returns:
+            ArtifactBuilder: The artifact builder object
         """
         if (rarity := self._artifact.rarity) > 0:
-            # Constraint: The level cannot exceed the max level for the artifact rarity
-            max_level = rarity * UPGRADE_STEP
-            expected_range = range(0, max_level + 1)
+            expected_range = range(0, self._artifact.max_level + 1)
             if level not in expected_range:
                 raise ValueError(
                     f"Invalid level '{level}' for rarity '{rarity}'. "
                     f"(Expected {min(expected_range)}-{max(expected_range)})"
                 )
 
-            # Constraint: The number of substats must be less than the rarity.
             if (substat_length := len(self._artifact.substats)) >= rarity:
                 raise ValueError(
                     f"Substat length mismatch with rarity '{rarity}' "
                     f"(Expected {rarity} substats, got {substat_length})"
                 )
-            self._artifact.mainstat.set_value_by_level(level)
+            if self._artifact.mainstat is not None:
+                self._artifact.mainstat.set_value_by_level(level)
 
         self._artifact.level = level
         return self
 
     def with_rarity(self, rarity: int) -> "ArtifactBuilder":
-        """Set the rarity of the artifact. Set any values that are dependent on the
-        rarity.
-
-        :param rarity: The rarity to set.
-        :type rarity: int
-        :raises ValueError: If the rarity is not within the valid range
-        :raises ValueError: If the number of substats exceeds the rarity of the artifact
-        :raises ValueError: If the substats are already full
-        :return: The artifact builder object
-        :rtype: ArtifactBuilder
+        """Set the rarity of the artifact. The rarity determines the number of substats.
+
+        Args:
+            rarity (int): The rarity to set.
+
+        Raises:
+            ValueError: If the number of substats exceeds the rarity of the artifact
+            ValueError: If the rarity is not within the valid range
+            ValueError: If the substats are already full
+
+        Returns:
+            ArtifactBuilder: The artifact builder object
         """
         if (level := self._artifact.level) > 0:
-            # Constraint: The level cannot exceed the max level for the artifact rarity
-            max_level = rarity * UPGRADE_STEP
+            max_level = rarity * UPGRADE_STEP if rarity > 2 else UPGRADE_STEP
             expected_range = range(0, max_level + 1)
             if level not in expected_range:
                 raise ValueError(
@@ -163,32 +169,41 @@ def with_rarity(self, rarity: int) -> "ArtifactBuilder":
         self._artifact.rarity = rarity
         return self
 
-    def with_set(self, artifact_set: str) -> "ArtifactBuilder":
+    def with_set(self, artifact_set: ArtifactSet) -> "ArtifactBuilder":
         """Set the artifact set.
 
-        :param set: The set to set.
-        :type set: str
-        :return: The artifact builder object
-        :rtype: ArtifactBuilder
+        Args:
+            artifact_set (ArtifactSet): The artifact set to set.
+
+        Returns:
+            ArtifactBuilder: The artifact builder object
         """
+        set_data = artifact_set.value
+        if self._artifact.artifact_slot is not None:
+            if self._artifact.artifact_slot not in set_data.slots:
+                raise ValueError(
+                    f"Invalid slot '{self._artifact.artifact_slot}' for set "
+                    f"'{artifact_set}' (expected: {set_data.slots})"
+                )
         self._artifact.artifact_set = artifact_set
         return self
 
     def with_slot(self, artifact_slot: ArtifactSlot) -> "ArtifactBuilder":
         """Set the artifact slot.
 
-        :param slot: The slot to set.
-        :type slot: str
-        :return: The artifact builder object
-        :rtype: ArtifactBuilder
+        Args:
+            artifact_slot (ArtifactSlot): The artifact slot to set.
+
+        Returns:
+            ArtifactBuilder: The artifact builder object
         """
         self._artifact.artifact_slot = artifact_slot
         return self
 
     def build(self) -> Artifact:
         """Build the artifact object based on the parameters passed into the builder.
 
-        :return: The built artifact object
-        :rtype: Artifact
+        Returns:
+            Artifact: The artifact object
         """
         return self._artifact

src/artipy/artifacts/upgrade_strategy.py

@@ -39,7 +39,8 @@ def upgrade(self, artifact: "Artifact") -> None:
         """
         new_level = artifact.level + 1
         artifact.level = new_level
-        artifact.mainstat.set_value_by_level(new_level)
+        if artifact.mainstat is not None:
+            artifact.mainstat.set_value_by_level(new_level)
 
 
 class AddStatStrategy(UpgradeStrategy):
@@ -58,7 +59,9 @@ def pick_stat(self, artifact: "Artifact") -> SubStat:
         :return: The new substat to add to the artifact.
         :rtype: SubStat
         """
-        stats = [s.name for s in (artifact.mainstat, *artifact.substats)]
+        stats = [
+            s.name for s in (artifact.mainstat, *artifact.substats) if s is not None
+        ]
         pool = {s: w for s, w in substat_weights.items() if s not in stats}
         population, weights = map(tuple, zip(*pool.items()))
         new_stat_name = choose(population, weights)