improve name to sku and vice versa

Author: offish

src/tf2_utils/__init__.py

@@ -1,12 +1,13 @@
 __title__ = "tf2-utils"
 __author__ = "offish"
-__version__ = "2.3.2"
+__version__ = "2.3.3"
 __license__ = "MIT"
 
 from .currency import CurrencyExchange
 from .exceptions import InvalidInventory, TF2UtilsError
 from .inventory import Inventory, map_inventory
 from .item import Item
+from .item_name import *
 from .marketplace_tf import (
     MarketplaceTF,
     MarketplaceTFException,
@@ -27,4 +28,4 @@
 from .sku import *
 from .utils import *
 
-# flake8: noqa
+# flake8: noqa: F401, F403

src/tf2_utils/item_name.py

@@ -0,0 +1,94 @@
+from tf2_data import EFFECTS, QUALITIES
+
+__all__ = [
+    "has_festivized_in_name",
+    "has_uncraftable_in_name",
+    "has_non_craftable_in_name",
+    "is_craftable",
+    "has_australium_in_name",
+    "has_strange_in_name",
+    "has_basic_killstreak_in_name",
+    "has_specialized_killstreak_in_name",
+    "has_professional_killstreak_in_name",
+    "has_killstreak_in_name",
+    "get_effect_in_name",
+    "get_quality_from_name",
+]
+
+QUALITY_NAMES = [
+    "Genuine",
+    "Vintage",
+    "Unusual",
+    "Unique",
+    "Strange",
+    "Haunted",
+    "Collector's",
+    "Decorated Weapon",
+]
+
+EFFECT_NAMES = [i for i in EFFECTS.keys() if not i.isnumeric()]
+
+
+def has_festivized_in_name(name: str) -> bool:
+    return "Festivized " in name
+
+
+def has_uncraftable_in_name(name: str) -> bool:
+    return "Uncraftable " in name
+
+
+def has_non_craftable_in_name(name: str) -> bool:
+    return "Non-Craftable " in name
+
+
+def is_craftable(name: str) -> bool:
+    return not (has_uncraftable_in_name(name) or has_non_craftable_in_name(name))
+
+
+def has_australium_in_name(name: str) -> bool:
+    return "Australium " in name
+
+
+def has_strange_in_name(name: str) -> bool:
+    return "Strange " in name
+
+
+def has_basic_killstreak_in_name(name: str) -> bool:
+    return name.startswith("Basic Killstreak ")
+
+
+def has_specialized_killstreak_in_name(name: str) -> bool:
+    return name.startswith("Specialized ")
+
+
+def has_professional_killstreak_in_name(name: str) -> bool:
+    return name.startswith("Professional ")
+
+
+def has_killstreak_in_name(name: str) -> bool:
+    return (
+        has_basic_killstreak_in_name(name)
+        or has_specialized_killstreak_in_name(name)
+        or has_professional_killstreak_in_name(name)
+    )
+
+
+def get_effect_in_name(name: str) -> int:
+    for effect in EFFECT_NAMES:
+        if effect in name:
+            return EFFECTS[effect]
+
+    return -1
+
+
+def get_quality_from_name(name: str) -> int:
+    quality = 6
+
+    for part in name.split(" "):
+        if part not in QUALITY_NAMES:
+            continue
+
+        quality = QUALITIES[part]
+        break
+
+    return quality

src/tf2_utils/schema.py

@@ -1,9 +1,22 @@
-from tf2_data import EFFECTS, KILLSTREAKS, QUALITIES, SchemaItems
+from tf2_data import SchemaItems
 from tf2_sku import to_sku
 
+from .item_name import (
+    get_effect_in_name,
+    get_quality_from_name,
+    has_australium_in_name,
+    has_basic_killstreak_in_name,
+    has_festivized_in_name,
+    has_professional_killstreak_in_name,
+    has_specialized_killstreak_in_name,
+    has_strange_in_name,
+    is_craftable,
+)
 from .sku import (
-    get_sku_effect,
-    get_sku_killstreak,
+    australium_in_sku,
+    festive_in_sku,
+    get_effect_name_from_sku,
+    get_killstreak_name_from_sku,
     sku_is_uncraftable,
     sku_to_defindex,
     sku_to_quality_name,
@@ -32,7 +45,7 @@ def defindex_to_full_name(self, defindex: int) -> str:
 
         return self.defindex_full_names.get(str(defindex), "")
 
-    def name_to_defindex(self, name: str) -> int:
+    def name_to_defindex(self, name: str, index: int = 0) -> int:
         if name == "Random Craft Weapon":
             return -50
 
@@ -57,9 +70,11 @@ def name_to_defindex(self, name: str) -> int:
         if name == "Name Tag":
             return defindexes[last_index]
 
-        return defindexes[0]
+        print(index)
+        print(defindexes)
+        return defindexes[index]
 
-    def defindex_to_image_url(self, defindex: int, large: bool = False) -> str:
+    def defindex_to_image_url(self, defindex: int, large_image: bool = False) -> str:
         # random craft weapon => shotgun
         if defindex == -50:
             defindex = 9
@@ -72,42 +87,28 @@ def defindex_to_image_url(self, defindex: int, large: bool = False) -> str:
             if item["defindex"] != defindex:
                 continue
 
-            return item["image_url"] if not large else item["image_url_large"]
+            return item["image_url_large"] if large_image else item["image_url"]
 
         return ""
 
-    def sku_to_image_url(self, sku: str, large: bool = False) -> str:
+    def sku_to_image_url(self, sku: str, large_image: bool = False) -> str:
         defindex = sku_to_defindex(sku)
-        return self.defindex_to_image_url(defindex, large)
-
-    def name_to_sku(self, name: str) -> str:
-        """This is not accurate, be careful when using this."""
-        parts = name.split(" ")
-
-        defindex = -1
-        craftable = True
-        quality = 6
+        return self.defindex_to_image_url(defindex, large_image)
 
-        for part in parts:
-            if part in ["Uncraftable", "Non-Craftable"]:
-                craftable = False
-
-            if part in QUALITIES:
-                quality = QUALITIES[part]
+    def get_defindex_from_name(self, name: str, entry_index: int = 0) -> int:
+        # try whole name, then remove everything till the
+        # first space for each iteration if defindex
+        # for that name doesnt exist
 
+        # example:
+        # Uncraftable Strange Team Captain => -1
+        # Strange Team Captain => -1
+        # Team Captain => 378 != -1, so break
         defindex_name = name
+        defindex = -1
 
         while True:
-            # try whole name, then remove everything till the
-            # first space for each iteration if defindex
-            # for that name doesnt exist
-
-            # example:
-            # Uncraftable Strange Team Captain => -1
-            # Strange Team Captain => -1
-            # Team Captain => 378 != -1, so break
-
-            defindex = self.name_to_defindex(defindex_name)
+            defindex = self.name_to_defindex(defindex_name, entry_index)
 
             if defindex != -1:
                 break
@@ -119,10 +120,46 @@ def name_to_sku(self, name: str) -> str:
 
             defindex_name = defindex_name[index + 1 :]
 
+        return defindex
+
+    def name_to_sku(self, name: str) -> str:
+        quality = get_quality_from_name(name)
+        defindex = self.get_defindex_from_name(name)
+        effect = get_effect_in_name(name)
+        is_australium = False
+        killstreak_tier = -1
+
+        if effect != -1:
+            quality = 5
+
+        if has_professional_killstreak_in_name(name):
+            killstreak_tier = 3
+
+        if has_specialized_killstreak_in_name(name):
+            killstreak_tier = 2
+
+        if has_basic_killstreak_in_name(name):
+            killstreak_tier = 1
+
+        # must be strange to be australium
+        if "Australium Gold" not in name and has_australium_in_name(name):
+            quality = 11
+            is_australium = True
+
+        # australium weapons are the second entry in the defindex list, need to
+        # get the defindex again
+        if is_australium:
+            defindex = self.get_defindex_from_name(name, 1)
+
         sku_properties = {
             "defindex": defindex,
             "quality": quality,
-            "craftable": craftable,
+            "effect": effect,
+            "strange": has_strange_in_name(name) and quality != 11,
+            "festivized": has_festivized_in_name(name),
+            "craftable": is_craftable(name),
+            "killstreak_tier": killstreak_tier,
+            "australium": is_australium,
         }
 
         return to_sku(sku_properties)
@@ -135,49 +172,54 @@ def sku_to_full_name(self, sku: str) -> str:
         defindex = sku_to_defindex(sku)
         return self.defindex_to_full_name(defindex)
 
-    def sku_to_name(self, sku: str, use_uncraftable: bool = True) -> str:
+    def format_name(
+        self,
+        item_name: str,
+        quality: str = "",
+        craftable: str = "",
+        effect: str = "",
+        killstreak: str = "",
+        strange: str = "",
+        festivized: str = "",
+    ) -> str:
+        return "".join(
+            [killstreak, strange, effect, craftable, quality, festivized, item_name]
+        )
+
+    def sku_to_name(self, sku: str, as_uncraftable: bool = True) -> str:
         name = self.sku_to_base_name(sku)
-        craftable = ""
         quality = sku_to_quality_name(sku)
-        effect = get_sku_effect(sku)
-        killstreak = get_sku_killstreak(sku)
-        strange = strange_in_sku(sku)
+        craftable = ""
+        festivized = "Festivized " if festive_in_sku(sku) else ""
+        effect = get_effect_name_from_sku(sku)
+        killstreak = get_killstreak_name_from_sku(sku)
+        is_australium = australium_in_sku(sku)
+        strange = ""
+
+        if strange_in_sku(sku):
+            strange = "Strange "
 
         if quality not in ["Unusual", "Unique"]:
             quality += " "
         else:
             quality = ""
 
-        if effect != -1:
-            effect = EFFECTS[str(effect)] + " "
-        else:
-            effect = ""
-
-        if killstreak != -1:
-            killstreak = KILLSTREAKS[str(killstreak)] + " "
-        else:
-            killstreak = ""
-
-        if strange:
-            strange = "Strange "
-        else:
-            strange = ""
+        if is_australium:
+            quality = ""
+            name = "Australium " + name
 
         if sku_is_uncraftable(sku):
-            if use_uncraftable:
+            if as_uncraftable:
                 craftable = "Uncraftable "
             else:
                 craftable = "Non-Craftable "
 
-        festive = ""
-        festivized = ""
-        wear = ""
-
-        # TODO: add killstreaks and other properties (strange unusual etc.)
-        # festive
-        # festivzed
-        # wear
-
-        return "".join(
-            [killstreak, strange, festive, festivized, craftable, quality, name, wear]
+        return self.format_name(
+            name,
+            quality=quality,
+            craftable=craftable,
+            effect=effect,
+            killstreak=killstreak,
+            strange=strange,
+            festivized=festivized,
         )