From 599a8b0dd732a3abadbd1c3fc5624f2813a3de17 Mon Sep 17 00:00:00 2001
From: Dima Gerasimov <karlicoss@gmail.com>
Date: Fri, 15 Apr 2022 13:34:09 +0100
Subject: [PATCH] ZipPath: support hash, iterdir and proper / operator

---
 my/core/kompress.py         | 34 +++++++++++++++++++++++++---------
 tests/core/test_kompress.py | 11 ++++++++++-
 2 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/my/core/kompress.py b/my/core/kompress.py
index b8e7724..8725bcf 100644
--- a/my/core/kompress.py
+++ b/my/core/kompress.py
@@ -6,7 +6,7 @@ from __future__ import annotations
 import pathlib
 from pathlib import Path
 import sys
-from typing import Union, IO, Sequence, Any
+from typing import Union, IO, Sequence, Any, Iterator
 import io
 
 PathIsh = Union[Path, str]
@@ -139,18 +139,22 @@ class ZipPath(ZipPathBase):
     root: zipfile.ZipFile
 
     @property
-    def filename(self) -> str:
+    def filepath(self) -> Path:
         res = self.root.filename
         assert res is not None  # make mypy happy
-        return res
+        return Path(res)
+
+    @property
+    def subpath(self) -> Path:
+        return Path(self.at)
 
     def absolute(self) -> ZipPath:
-        return ZipPath(Path(self.filename).absolute(), self.at)
+        return ZipPath(self.filepath.absolute(), self.at)
 
     def exists(self) -> bool:
         if self.at == '':
             # special case, the base class returns False in this case for some reason
-            return Path(self.filename).exists()
+            return self.filepath.exists()
         return super().exists()
 
     def rglob(self, glob: str) -> Sequence[ZipPath]:
@@ -162,16 +166,25 @@ class ZipPath(ZipPathBase):
 
     def relative_to(self, other: ZipPath) -> Path:
         assert self.root == other.root, (self.root, other.root)
-        return Path(self.at).relative_to(Path(other.at))
+        return self.subpath.relative_to(other.subpath)
 
     @property
     def parts(self) -> Sequence[str]:
         # messy, but might be ok..
-        return Path(self.filename).parts + Path(self.at).parts
+        return self.filepath.parts + self.subpath.parts
+
+    def __truediv__(self, key) -> ZipPath:
+        # need to implement it so the return type is not zipfile.Path
+        s = super().__truediv__(key)
+        return ZipPath(s.root, s.at)  # type: ignore[attr-defined]
+
+    def iterdir(self) -> Iterator[ZipPath]:
+        for s in super().iterdir():
+            yield ZipPath(s.root, s.at)  # type: ignore[attr-defined]
 
     @property
     def stem(self) -> str:
-        return Path(self.at).stem
+        return self.subpath.stem
 
     @property  # type: ignore[misc]
     def __class__(self):
@@ -181,4 +194,7 @@ class ZipPath(ZipPathBase):
         # hmm, super class doesn't seem to treat as equals unless they are the same object
         if not isinstance(other, ZipPath):
             return False
-        return self.filename == other.filename and Path(self.at) == Path(other.at)
+        return (self.filepath, self.subpath) == (other.filepath, other.subpath)
+
+    def __hash__(self) -> int:
+        return hash((self.filepath, self.subpath))
diff --git a/tests/core/test_kompress.py b/tests/core/test_kompress.py
index 3561444..949a7f1 100644
--- a/tests/core/test_kompress.py
+++ b/tests/core/test_kompress.py
@@ -63,12 +63,17 @@ def test_zippath() -> None:
 
     # magic! convenient to make third party libraries agnostic of ZipPath
     assert isinstance(zp, Path)
+    assert isinstance(zp, ZipPath)
+    assert isinstance(zp / 'subpath', Path)
     # TODO maybe change __str__/__repr__? since it's a bit misleading:
     # Path('/code/hpi/tests/core/structure_data/gdpr_export.zip', 'gdpr_export/')
 
     assert ZipPath(target) == ZipPath(target)
     assert zp.absolute() == zp
 
+    # shouldn't crash
+    hash(zp)
+
     assert zp.exists()
     assert (zp / 'gdpr_export/comments').exists()
     # check str constructor just in case
@@ -77,7 +82,7 @@ def test_zippath() -> None:
 
     matched = list(zp.rglob('*'))
     assert len(matched) > 0
-    assert all(p.filename == str(target) for p in matched), matched
+    assert all(p.filepath == target for p in matched), matched
 
     rpaths = [str(p.relative_to(zp)) for p in matched]
     assert rpaths == [
@@ -106,3 +111,7 @@ def test_zippath() -> None:
     ]
 
     assert list(zp.rglob('mes*')) == [ZipPath(target, 'gdpr_export/messages')]
+
+    iterdir_res = list((zp / 'gdpr_export').iterdir())
+    assert len(iterdir_res) == 3
+    assert all(isinstance(p, Path) for p in iterdir_res)