Add path type __eq__ and __repr__ QOL changes (#79)

Co-authored-by: Yun Zheng Hu <[email protected]>
fox-it · Aug 25, 2023 · b2818b0 · b2818b0
1 parent ddea907
commit b2818b0
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 11 deletions.
diff --git a/flow/record/fieldtypes/__init__.py b/flow/record/fieldtypes/__init__.py
@@ -643,6 +643,14 @@ def __new__(cls, *args):
 
         return cls._from_parts(args)
 
+    def __eq__(self, other: Any) -> bool:
+        if isinstance(other, str):
+            return str(self) == other or self == self.__class__(other)
+        return super().__eq__(other)
+
+    def __repr__(self) -> str:
+        return repr(str(self))
+
     def _pack(self):
         path_type = PATH_WINDOWS if isinstance(self, windows_path) else PATH_POSIX
         return (str(self), path_type)
@@ -674,4 +682,13 @@ class posix_path(pathlib.PurePosixPath, path):
 
 
 class windows_path(pathlib.PureWindowsPath, path):
-    pass
+    def __repr__(self) -> str:
+        s = str(self)
+        quote = "'"
+        if "'" in s:
+            if '"' in s:
+                s = s.replace("'", "\\'")
+            else:
+                quote = '"'
+
+        return f"{quote}{s}{quote}"
diff --git a/tests/test_fieldtypes.py b/tests/test_fieldtypes.py
@@ -16,7 +16,7 @@
     _is_windowslike_path,
 )
 from flow.record.fieldtypes import datetime as dt
-from flow.record.fieldtypes import fieldtype_for_value, net, uri
+from flow.record.fieldtypes import fieldtype_for_value, net, uri, windows_path
 
 UTC = timezone.utc
 
@@ -588,6 +588,7 @@ def test_path():
 
     r = TestRecord("")
     assert str(r.value) == "."
+    assert r.value == "."
 
     if os.name == "nt":
         native_path_str = windows_path_str
@@ -693,7 +694,7 @@ def test_path_posix(path_initializer, path, expected_repr):
     )
 
     record = TestRecord(path=path_initializer(path))
-    assert repr(record) == f"<test/path path=posix_path('{expected_repr}')>"
+    assert repr(record) == f"<test/path path='{expected_repr}'>"
 
 
 @pytest.mark.parametrize(
@@ -707,20 +708,27 @@ def test_path_posix(path_initializer, path, expected_repr):
 @pytest.mark.parametrize(
     "path,expected_repr,expected_str",
     [
-        ("c:\\windows\\temp\\foo\\bar", "c:/windows/temp/foo/bar", r"c:\windows\temp\foo\bar"),
-        (r"C:\Windows\Temp\foo\bar", "C:/Windows/Temp/foo/bar", r"C:\Windows\Temp\foo\bar"),
-        (r"d:/Users/Public", "d:/Users/Public", r"d:\Users\Public"),
+        ("c:\\windows\\temp\\foo\\bar", r"'c:\windows\temp\foo\bar'", r"c:\windows\temp\foo\bar"),
+        (r"C:\Windows\Temp\foo\bar", r"'C:\Windows\Temp\foo\bar'", r"C:\Windows\Temp\foo\bar"),
+        (r"d:/Users/Public", r"'d:\Users\Public'", r"d:\Users\Public"),
         (
             "/sysvol/Windows/System32/drivers/null.sys",
-            "/sysvol/Windows/System32/drivers/null.sys",
+            r"'\sysvol\Windows\System32\drivers\null.sys'",
             r"\sysvol\Windows\System32\drivers\null.sys",
         ),
         (
             "/c:/Windows/System32/drivers/null.sys",
-            "/c:/Windows/System32/drivers/null.sys",
+            r"'\c:\Windows\System32\drivers\null.sys'",
             r"\c:\Windows\System32\drivers\null.sys",
         ),
-        ("Users\\Public", "Users/Public", r"Users\Public"),
+        ("Users\\Public", r"'Users\Public'", r"Users\Public"),
+        (r"i:\don't.exe", '"i:\\don\'t.exe"', r"i:\don't.exe"),
+        (
+            'y:\\shakespeare\\"to be or not to be".txt',
+            "'y:\\shakespeare\\\"to be or not to be\".txt'",
+            'y:\\shakespeare\\"to be or not to be".txt',
+        ),
+        ("c:\\my'quotes\".txt", "'c:\\my\\'quotes\".txt'", "c:\\my'quotes\".txt"),
     ],
 )
 def test_path_windows(path_initializer, path, expected_repr, expected_str):
@@ -731,10 +739,20 @@ def test_path_windows(path_initializer, path, expected_repr, expected_str):
         ],
     )
     record = TestRecord(path=path_initializer(path))
-    assert repr(record) == f"<test/path path=windows_path('{expected_repr}')>"
+    assert repr(record) == f"<test/path path={expected_repr}>"
+    assert repr(record.path) == expected_repr
     assert str(record.path) == expected_str
 
 
+def test_windows_path_eq():
+    path = windows_path("c:\\windows\\test.exe")
+    assert path == "c:\\windows\\test.exe"
+    assert path == "c:/windows/test.exe"
+    assert path == "c:/windows\\test.exe"
+    assert path == "c:\\WINDOWS\\tEsT.ExE"
+    assert path != "c:/windows\\test2.exe"
+
+
 def test_fieldtype_for_value():
     assert fieldtype_for_value(True) == "boolean"
     assert fieldtype_for_value(False) == "boolean"

diff --git a/tests/test_regression.py b/tests/test_regression.py
@@ -536,7 +536,7 @@ def test_windows_path_regression(path_initializer):
     )
     r = TestRecord(path=path_initializer("/c:/Windows/System32/drivers/null.sys"))
     assert str(r.path) == "\\c:\\Windows\\System32\\drivers\\null.sys"
-    assert repr(r.path) == "windows_path('/c:/Windows/System32/drivers/null.sys')"
+    assert repr(r.path) == "'\\c:\\Windows\\System32\\drivers\\null.sys'"
 
 
 @pytest.mark.parametrize(