improve types according to test/test_tree.py

Author: BenediktSeidl

pygit2/_pygit2.pyi

@@ -1117,11 +1117,11 @@ class Tree(Object):
         interhunk_lines: int = 0,
     ) -> Diff: ...
     def __contains__(self, other: str) -> bool: ...  # Tree_contains
-    def __getitem__(self, index: str | int) -> Object: ...  # Tree_subscript
+    def __getitem__(self, index: str | int) -> Tree | Blob: ...  # Tree_subscript
     def __iter__(self) -> Iterator[Object]: ...
     def __len__(self) -> int: ...  # Tree_len
-    def __rtruediv__(self, other: str) -> Object: ...
-    def __truediv__(self, other: str) -> Object: ...  # Tree_divide
+    def __rtruediv__(self, other: str) -> Tree | Blob: ...
+    def __truediv__(self, other: str) -> Tree | Blob: ...  # Tree_divide
 
 class TreeBuilder:
     def clear(self) -> None: ...

test/test_tree.py

@@ -28,6 +28,7 @@
 import pytest
 
 import pygit2
+from pygit2 import Object, Repository, Tree
 from pygit2.enums import FileMode, ObjectType
 
 from . import utils
@@ -36,19 +37,20 @@
 SUBTREE_SHA = '614fd9a3094bf618ea938fffc00e7d1a54f89ad0'
 
 
-def assertTreeEntryEqual(entry, sha, name, filemode):
+def assertTreeEntryEqual(entry: Object, sha: str, name: str, filemode: int) -> None:
     assert entry.id == sha
     assert entry.name == name
     assert entry.filemode == filemode
     assert entry.raw_name == name.encode('utf-8')
 
 
-def test_read_tree(barerepo):
+def test_read_tree(barerepo: Repository) -> None:
     tree = barerepo[TREE_SHA]
+    assert isinstance(tree, Tree)
     with pytest.raises(TypeError):
-        tree[()]
+        tree[()]  # type: ignore
     with pytest.raises(TypeError):
-        tree / 123
+        tree / 123  # type: ignore
     utils.assertRaisesWithArg(KeyError, 'abcd', lambda: tree['abcd'])
     utils.assertRaisesWithArg(IndexError, -4, lambda: tree[-4])
     utils.assertRaisesWithArg(IndexError, 3, lambda: tree[3])
@@ -72,45 +74,50 @@ def test_read_tree(barerepo):
     sha = '297efb891a47de80be0cfe9c639e4b8c9b450989'
     assertTreeEntryEqual(tree['c/d'], sha, 'd', 0o0100644)
     assertTreeEntryEqual(tree / 'c/d', sha, 'd', 0o0100644)
-    assertTreeEntryEqual(tree / 'c' / 'd', sha, 'd', 0o0100644)
-    assertTreeEntryEqual(tree['c']['d'], sha, 'd', 0o0100644)
-    assertTreeEntryEqual((tree / 'c')['d'], sha, 'd', 0o0100644)
+    assertTreeEntryEqual(tree / 'c' / 'd', sha, 'd', 0o0100644)  # type: ignore[operator]
+    assertTreeEntryEqual(tree['c']['d'], sha, 'd', 0o0100644)  # type: ignore[index]
+    assertTreeEntryEqual((tree / 'c')['d'], sha, 'd', 0o0100644)  # type: ignore[index]
     utils.assertRaisesWithArg(KeyError, 'ab/cd', lambda: tree['ab/cd'])
     utils.assertRaisesWithArg(KeyError, 'ab/cd', lambda: tree / 'ab/cd')
-    utils.assertRaisesWithArg(KeyError, 'ab', lambda: tree / 'c' / 'ab')
+    utils.assertRaisesWithArg(KeyError, 'ab', lambda: tree / 'c' / 'ab')  # type: ignore[operator]
     with pytest.raises(TypeError):
-        tree / 'a' / 'cd'
+        tree / 'a' / 'cd'  # type: ignore
 
 
-def test_equality(barerepo):
+def test_equality(barerepo: Repository) -> None:
     tree_a = barerepo['18e2d2e9db075f9eb43bcb2daa65a2867d29a15e']
     tree_b = barerepo['2ad1d3456c5c4a1c9e40aeeddb9cd20b409623c8']
+    assert isinstance(tree_a, Tree)
+    assert isinstance(tree_b, Tree)
 
     assert tree_a['a'] != tree_b['a']
     assert tree_a['a'] != tree_b['b']
     assert tree_a['b'] == tree_b['b']
 
 
-def test_sorting(barerepo):
+def test_sorting(barerepo: Repository) -> None:
     tree_a = barerepo['18e2d2e9db075f9eb43bcb2daa65a2867d29a15e']
+    assert isinstance(tree_a, Tree)
     assert list(tree_a) == sorted(reversed(list(tree_a)), key=pygit2.tree_entry_key)
-    assert list(tree_a) != reversed(list(tree_a))
+    assert list(tree_a) != reversed(list(tree_a))  # type: ignore[comparison-overlap]
 
 
-def test_read_subtree(barerepo):
+def test_read_subtree(barerepo: Repository) -> None:
     tree = barerepo[TREE_SHA]
+    assert isinstance(tree, Tree)
 
     subtree_entry = tree['c']
     assertTreeEntryEqual(subtree_entry, SUBTREE_SHA, 'c', 0o0040000)
-    assert subtree_entry.type == ObjectType.TREE
+    assert subtree_entry.type == int(ObjectType.TREE)
     assert subtree_entry.type_str == 'tree'
 
     subtree_entry = tree / 'c'
     assertTreeEntryEqual(subtree_entry, SUBTREE_SHA, 'c', 0o0040000)
-    assert subtree_entry.type == ObjectType.TREE
+    assert subtree_entry.type == int(ObjectType.TREE)
     assert subtree_entry.type_str == 'tree'
 
     subtree = barerepo[subtree_entry.id]
+    assert isinstance(subtree, Tree)
     assert 1 == len(subtree)
     sha = '297efb891a47de80be0cfe9c639e4b8c9b450989'
     assertTreeEntryEqual(subtree[0], sha, 'd', 0o0100644)
@@ -119,7 +126,7 @@ def test_read_subtree(barerepo):
     assert subtree_entry == barerepo[subtree_entry.id]
 
 
-def test_new_tree(barerepo):
+def test_new_tree(barerepo: Repository) -> None:
     repo = barerepo
     b0 = repo.create_blob('1')
     b1 = repo.create_blob('2')
@@ -138,8 +145,8 @@ def test_new_tree(barerepo):
         ('y', b1, pygit2.Blob, FileMode.BLOB_EXECUTABLE, ObjectType.BLOB, 'blob'),
         ('z', subtree.id, pygit2.Tree, FileMode.TREE, ObjectType.TREE, 'tree'),
     ]:
-        assert name in tree
-        obj = tree[name]
+        assert name in tree  # type: ignore[operator]
+        obj = tree[name]  # type: ignore[index]
         assert isinstance(obj, cls)
         assert obj.name == name
         assert obj.filemode == filemode
@@ -148,7 +155,7 @@ def test_new_tree(barerepo):
         assert repo[obj.id].id == oid
         assert obj == repo[obj.id]
 
-        obj = tree / name
+        obj = tree / name  # type: ignore[operator]
         assert isinstance(obj, cls)
         assert obj.name == name
         assert obj.filemode == filemode
@@ -158,44 +165,49 @@ def test_new_tree(barerepo):
         assert obj == repo[obj.id]
 
 
-def test_modify_tree(barerepo):
+def test_modify_tree(barerepo: Repository) -> None:
     tree = barerepo[TREE_SHA]
     with pytest.raises(TypeError):
-        operator.setitem('c', tree['a'])
+        operator.setitem('c', tree['a'])  # type: ignore
     with pytest.raises(TypeError):
-        operator.delitem('c')
+        operator.delitem('c')  # type: ignore
 
 
-def test_iterate_tree(barerepo):
+def test_iterate_tree(barerepo: Repository) -> None:
     """
     Testing that we're able to iterate of a Tree object and that the
     resulting sha strings are consistent with the sha strings we could
     get with other Tree access methods.
     """
     tree = barerepo[TREE_SHA]
+    assert isinstance(tree, Tree)
     for tree_entry in tree:
+        assert tree_entry.name is not None
         assert tree_entry == tree[tree_entry.name]
 
 
-def test_iterate_tree_nested(barerepo):
+def test_iterate_tree_nested(barerepo: Repository) -> None:
     """
     Testing that we're able to iterate of a Tree object and then iterate
     trees we receive as a result.
     """
     tree = barerepo[TREE_SHA]
+    assert isinstance(tree, Tree)
     for tree_entry in tree:
         if isinstance(tree_entry, pygit2.Tree):
             for tree_entry2 in tree_entry:
                 pass
 
 
-def test_deep_contains(barerepo):
+def test_deep_contains(barerepo: Repository) -> None:
     tree = barerepo[TREE_SHA]
+    assert isinstance(tree, Tree)
     assert 'a' in tree
     assert 'c' in tree
     assert 'c/d' in tree
     assert 'c/e' not in tree
     assert 'd' not in tree
 
+    assert isinstance(tree['c'], Tree)
     assert 'd' in tree['c']
     assert 'e' not in tree['c']