fix(UnityVersion): handle optional postfixes, fixes #369

Author: K0lb3

UnityPy/helpers/UnityVersion.py

@@ -4,7 +4,10 @@
 from enum import IntEnum
 from typing import Optional, Tuple, Union
 
-VersionPattern = re.compile(r"^(?P<major>\d+)\.(?P<minor>\d+)\.(?P<build>\d+)(?P<type_str>.+?)?(?P<type_number>\d+)?$")
+VersionPattern = re.compile(
+    r"^(?P<major>\d+)\.(?P<minor>\d+)\.(?P<build>\d+)(?P<type_str>.+?)?(?P<type_number>\d+)?(?P<postfix>.*)$",
+    flags=re.DOTALL,
+)
 
 
 class UnityVersionType(IntEnum):
@@ -20,6 +23,7 @@ class UnityVersionType(IntEnum):
 class UnityVersion(int):
     # https://github.com/AssetRipper/VersionUtilities/blob/master/VersionUtilities/UnityVersion.cs
     _type_str: Optional[str]
+    _postfix: Optional[str]
 
     @property
     def major(self):
@@ -41,6 +45,10 @@ def type(self):
     def type_str(self):
         return getattr(self, "_type_str", self.type.name)
 
+    @property
+    def postfix(self):
+        return getattr(self, "_postfix", "")
+
     @property
     def type_number(self):
         return self & 0xFF
@@ -56,6 +64,7 @@ def from_str(cls, version: str):
         # formats:
         #   old: 5.0.0, <major>.<minor>.<build>
         #   new: 2018.1.1f2 <major>.<minor>.<build><type_str><type_number>
+        #   the new format string can be followed by a custom postfix
         match = VersionPattern.match(version)
         if not match:
             raise ValueError(f"Invalid version string: {version}")
@@ -64,6 +73,7 @@ def from_str(cls, version: str):
         build = int(match.group("build"))
         type_str = match.group("type_str")
         type_number = int(match.group("type_number") or 0)
+        postfix = match.group("postfix")
 
         if type_str is None:
             return cls.from_list(major, minor, build)
@@ -72,10 +82,19 @@ def from_str(cls, version: str):
         obj = cls.from_list(major, minor, build, type, type_number)
         if type is UnityVersionType.u:
             obj._type_str = type_str
+        if postfix:
+            obj._postfix = postfix
+
         return obj
 
+    def __str__(self) -> str:
+        if self.major <= 5:
+            return f"{self.major}.{self.minor}.{self.build}"
+        else:
+            return f"{self.major}.{self.minor}{self.type_str}{self.type_number}{self.postfix}"
+
     def __repr__(self) -> str:
-        return f"UnityVersion {self.major}.{self.minor}{self.type_str}{self.type_number}"
+        return f"UnityVersion {self.__str__()}"
 
     def __getitem__(self, idx: Union[int, slice]) -> Union[int, Tuple[int, ...]]:
         values = (

tests/test_UnityVersion.py

@@ -14,6 +14,7 @@
         ("2021.1.0c1", (2021, 1, 0, UnityVersionType.c.value, 1)),
         ("2022.2.0x1", (2022, 2, 0, UnityVersionType.x.value, 1)),
         ("2018.1.1z2", (2018, 1, 1, UnityVersionType.u.value, 2)),  # unknown type
+        ("2022.3.62f2\n2", (2022, 3, 62, UnityVersionType.f.value, 2)),
     ],
 )
 def test_parse_unity_version(version_str, expected_tuple):
@@ -24,6 +25,7 @@ def test_parse_unity_version(version_str, expected_tuple):
     assert v.build == expected_tuple[2]
     assert v.type.value == expected_tuple[3]
     assert v.type_number == expected_tuple[4]
+    assert UnityVersion.from_list(*expected_tuple) == v
 
 
 @pytest.mark.parametrize(
@@ -58,6 +60,7 @@ def test_comparison_with_tuple(version_str, compare_tuple):
         ("2018.1.1f2", "2018.1.1f1"),
         ("2018.1.1f2", "2018.1.2f2"),
         ("2018.1.1f2", "2018.2.1f2"),
+        ("2022.3.62f2\n2", "2022.3.62f2"),
     ],
 )
 def test_comparison_with_unityversion(version_str, other_str):
@@ -69,12 +72,3 @@ def test_comparison_with_unityversion(version_str, other_str):
     assert (v1 <= v2) == (v1.as_tuple() <= v2.as_tuple())
     assert (v1 > v2) == (v1.as_tuple() > v2.as_tuple())
     assert (v1 >= v2) == (v1.as_tuple() >= v2.as_tuple())
-
-
-def test_repr_and_str():
-    v = UnityVersion.from_str("2018.1.1f2")
-    assert "UnityVersion" in repr(v)
-    assert str(v.major) in repr(v)
-    assert str(v.minor) in repr(v)
-    assert v.type_str in repr(v)
-    assert str(v.type_number) in repr(v)