improve types according to test/test_index.py

Author: BenediktSeidl

pygit2/_libgit2/ffi.pyi

@@ -42,6 +42,11 @@ class _Pointer(Generic[T]):
     @overload
     def __getitem__(self, item: slice[None, None, None]) -> bytes: ...
 
+class ArrayC(Generic[T]):
+    # incomplete!
+    # def _len(self, ?) -> ?: ...
+    pass
+
 class GitTimeC:
     # incomplete
     time: int
@@ -132,13 +137,13 @@ class GitDescribeFormatOptionsC:
     version: int
     abbreviated_size: int
     always_use_long_format: int
-    dirty_suffix: char_pointer
+    dirty_suffix: ArrayC[char]
 
 class GitDescribeOptionsC:
     version: int
     max_candidates_tags: int
     describe_strategy: int
-    pattern: char_pointer
+    pattern: ArrayC[char]
     only_follow_first_parent: int
     show_commit_oid_as_fallback: int
 
@@ -148,6 +153,11 @@ class GitDescribeResultC:
 class GitIndexC:
     pass
 
+class GitIndexEntryC:
+    # incomplete?
+    mode: int
+    path: ArrayC[char]
+
 class GitMergeFileResultC:
     pass
 
@@ -158,11 +168,13 @@ class GitStashSaveOptionsC:
     version: int
     flags: int
     stasher: GitSignatureC
-    message: char_pointer
+    message: ArrayC[char]
     paths: GitStrrayC
 
 class GitStrrayC:
-    pass
+    # incomplete?
+    strings: NULL_TYPE | ArrayC[char]
+    count: int
 
 class GitTreeC:
     pass
@@ -171,11 +183,11 @@ class GitRepositoryInitOptionsC:
     version: int
     flags: int
     mode: int
-    workdir_path: char_pointer
-    description: char_pointer
-    template_path: char_pointer
-    initial_head: char_pointer
-    origin_url: char_pointer
+    workdir_path: ArrayC[char]
+    description: ArrayC[char]
+    template_path: ArrayC[char]
+    initial_head: ArrayC[char]
+    origin_url: ArrayC[char]
 
 class GitCloneOptionsC:
     pass
@@ -229,6 +241,10 @@ def new(a: Literal['git_attr_options *']) -> GitAttrOptionsC: ...
 @overload
 def new(a: Literal['git_buf *']) -> GitBufC: ...
 @overload
+def new(a: Literal['char *'], b: bytes) -> char_pointer: ...
+@overload
+def new(a: Literal['char *[]'], b: list[char_pointer]) -> ArrayC[char_pointer]: ...
+@overload
 def new(a: Literal['git_checkout_options *']) -> GitCheckoutOptionsC: ...
 @overload
 def new(a: Literal['git_commit **']) -> _Pointer[GitCommitC]: ...
@@ -251,6 +267,8 @@ def new(a: Literal['struct git_reference **']) -> _Pointer[GitReferenceC]: ...
 @overload
 def new(a: Literal['git_index **']) -> _Pointer[GitIndexC]: ...
 @overload
+def new(a: Literal['git_index_entry *']) -> GitIndexEntryC: ...
+@overload
 def new(a: Literal['git_merge_file_result *']) -> GitMergeFileResultC: ...
 @overload
 def new(a: Literal['git_object *']) -> GitObjectC: ...
@@ -263,13 +281,15 @@ def new(a: Literal['git_signature **']) -> _Pointer[GitSignatureC]: ...
 @overload
 def new(a: Literal['git_stash_save_options *']) -> GitStashSaveOptionsC: ...
 @overload
+def new(a: Literal['git_strarray *']) -> GitStrrayC: ...
+@overload
 def new(a: Literal['git_tree **']) -> _Pointer[GitTreeC]: ...
 @overload
 def new(a: Literal['git_buf *'], b: tuple[NULL_TYPE, Literal[0]]) -> GitBufC: ...
 @overload
 def new(a: Literal['char **']) -> _Pointer[char_pointer]: ...
 @overload
-def new(a: Literal['char[]', 'char []'], b: bytes | NULL_TYPE) -> char_pointer: ...
+def new(a: Literal['char[]', 'char []'], b: bytes | NULL_TYPE) -> ArrayC[char]: ...
 def addressof(a: object, attribute: str) -> _Pointer[object]: ...
 
 class buffer(bytes):

pygit2/_pygit2.pyi

@@ -723,6 +723,7 @@ class Branches:
 
 class Repository:
     _pointer: GitRepositoryC
+    _repo: GitRepositoryC
     default_signature: Signature
     head: Reference
     head_is_detached: bool
@@ -788,7 +789,7 @@ class Repository:
     def compress_references(self) -> None: ...
     @property
     def config(self) -> Config: ...
-    def create_blob(self, data: bytes) -> Oid: ...
+    def create_blob(self, data: str | bytes) -> Oid: ...
     def create_blob_fromdisk(self, path: str) -> Oid: ...
     def create_blob_fromiobase(self, iobase: IOBase) -> Oid: ...
     def create_blob_fromworkdir(self, path: str | Path) -> Oid: ...

pygit2/index.py

@@ -26,9 +26,10 @@
 import typing
 import warnings
 from dataclasses import dataclass
+from os import PathLike
 
 # Import from pygit2
-from ._pygit2 import Diff, Oid, Tree
+from ._pygit2 import Diff, Oid, Repository, Tree
 from .enums import DiffOption, FileMode
 from .errors import check_error
 from .ffi import C, ffi
@@ -41,7 +42,7 @@ class Index:
     # a proper implementation in some places: e.g. checking the index type
     # from C code (see Tree_diff_to_index)
 
-    def __init__(self, path: str | None = None) -> None:
+    def __init__(self, path: str | PathLike[str] | None = None) -> None:
         """Create a new Index
 
         If path is supplied, the read and write methods will use that path
@@ -116,16 +117,16 @@ def read(self, force=True):
         err = C.git_index_read(self._index, force)
         check_error(err, io=True)
 
-    def write(self):
+    def write(self) -> None:
         """Write the contents of the Index to disk."""
         err = C.git_index_write(self._index)
         check_error(err, io=True)
 
-    def clear(self):
+    def clear(self) -> None:
         err = C.git_index_clear(self._index)
         check_error(err)
 
-    def read_tree(self, tree):
+    def read_tree(self, tree: Oid | Tree | str) -> None:
         """Replace the contents of the Index with those of the given tree,
         expressed either as a <Tree> object or as an oid (string or <Oid>).
 
@@ -134,6 +135,8 @@ def read_tree(self, tree):
         """
         repo = self._repo
         if isinstance(tree, str):
+            if repo is None:
+                raise TypeError('id given but no associated repository')
             tree = repo[tree]
 
         if isinstance(tree, Oid):
@@ -142,14 +145,14 @@ def read_tree(self, tree):
 
             tree = repo[tree]
         elif not isinstance(tree, Tree):
-            raise TypeError('argument must be Oid or Tree')
+            raise TypeError('argument must be Oid, Tree or str')
 
         tree_cptr = ffi.new('git_tree **')
         ffi.buffer(tree_cptr)[:] = tree._pointer[:]
         err = C.git_index_read_tree(self._index, tree_cptr[0])
         check_error(err)
 
-    def write_tree(self, repo=None):
+    def write_tree(self, repo: Repository | None = None) -> Oid:
         """Create a tree out of the Index. Return the <Oid> object of the
         written tree.
 
@@ -172,23 +175,23 @@ def write_tree(self, repo=None):
         check_error(err)
         return Oid(raw=bytes(ffi.buffer(coid)[:]))
 
-    def remove(self, path, level=0):
+    def remove(self, path: PathLike[str] | str, level: int = 0) -> None:
         """Remove an entry from the Index."""
         err = C.git_index_remove(self._index, to_bytes(path), level)
         check_error(err, io=True)
 
-    def remove_directory(self, path, level=0):
+    def remove_directory(self, path: PathLike[str] | str, level: int = 0) -> None:
         """Remove a directory from the Index."""
         err = C.git_index_remove_directory(self._index, to_bytes(path), level)
         check_error(err, io=True)
 
-    def remove_all(self, pathspecs):
+    def remove_all(self, pathspecs: typing.Sequence[str | PathLike[str]]) -> None:
         """Remove all index entries matching pathspecs."""
         with StrArray(pathspecs) as arr:
             err = C.git_index_remove_all(self._index, arr.ptr, ffi.NULL, ffi.NULL)
             check_error(err, io=True)
 
-    def add_all(self, pathspecs=None):
+    def add_all(self, pathspecs: None | list[str | PathLike[str]] = None) -> None:
         """Add or update index entries matching files in the working directory.
 
         If pathspecs are specified, only files matching those pathspecs will
@@ -199,7 +202,7 @@ def add_all(self, pathspecs=None):
             err = C.git_index_add_all(self._index, arr.ptr, 0, ffi.NULL, ffi.NULL)
             check_error(err, io=True)
 
-    def add(self, path_or_entry):
+    def add(self, path_or_entry: 'IndexEntry | str | PathLike[str]') -> None:
         """Add or update an entry in the Index.
 
         If a path is given, that file will be added. The path must be relative
@@ -217,11 +220,13 @@ def add(self, path_or_entry):
             path = path_or_entry
             err = C.git_index_add_bypath(self._index, to_bytes(path))
         else:
-            raise TypeError('argument must be string or IndexEntry')
+            raise TypeError('argument must be string, Path or IndexEntry')
 
         check_error(err, io=True)
 
-    def add_conflict(self, ancestor, ours, theirs):
+    def add_conflict(
+        self, ancestor: 'IndexEntry', ours: 'IndexEntry', theirs: 'IndexEntry | None'
+    ) -> None:
         """
         Add or update index entries to represent a conflict. Any staged entries that
         exist at the given paths will be removed.
@@ -243,7 +248,9 @@ def add_conflict(self, ancestor, ours, theirs):
         if theirs and not isinstance(theirs, IndexEntry):
             raise TypeError('theirs has to be an instance of IndexEntry or None')
 
-        centry_ancestor = centry_ours = centry_theirs = ffi.NULL
+        centry_ancestor: ffi.NULL_TYPE | ffi.GitIndexEntryC = ffi.NULL
+        centry_ours: ffi.NULL_TYPE | ffi.GitIndexEntryC = ffi.NULL
+        centry_theirs: ffi.NULL_TYPE | ffi.GitIndexEntryC = ffi.NULL
         if ancestor is not None:
             centry_ancestor, _ = ancestor._to_c()
         if ours is not None:
@@ -418,7 +425,7 @@ def _from_c(cls, centry):
 
 
 class IndexEntry:
-    path: str
+    path: str | PathLike[str]
     'The path of this entry'
 
     id: Oid
@@ -427,7 +434,9 @@ class IndexEntry:
     mode: FileMode
     'The mode of this entry, a FileMode value'
 
-    def __init__(self, path, object_id: Oid, mode: FileMode):
+    def __init__(
+        self, path: str | PathLike[str], object_id: Oid, mode: FileMode
+    ) -> None:
         self.path = path
         self.id = object_id
         self.mode = mode
@@ -459,7 +468,7 @@ def __eq__(self, other):
             self.path == other.path and self.id == other.id and self.mode == other.mode
         )
 
-    def _to_c(self):
+    def _to_c(self) -> tuple['ffi.GitIndexEntryC', 'ffi.ArrayC[ffi.char]']:
         """Convert this entry into the C structure
 
         The first returned arg is the pointer, the second is the reference to

pygit2/repository.py

@@ -1183,7 +1183,7 @@ def stash(
 
         if paths:
             arr = StrArray(paths)
-            opts.paths = arr.ptr[0]
+            opts.paths = arr.ptr[0]  # type: ignore[index]
 
         coid = ffi.new('git_oid *')
         err = C.git_stash_save_with_opts(coid, self._repo, opts)

pygit2/utils.py

@@ -25,11 +25,25 @@
 
 import contextlib
 import os
-from typing import Generic, Iterator, Protocol, TypeVar, Union, overload
+from types import TracebackType
+from typing import (
+    TYPE_CHECKING,
+    Generic,
+    Iterator,
+    Optional,
+    Protocol,
+    Sequence,
+    TypeVar,
+    Union,
+    overload,
+)
 
 # Import from pygit2
 from .ffi import C, ffi
 
+if TYPE_CHECKING:
+    from ._libgit2.ffi import ArrayC, GitStrrayC, char
+
 
 def maybe_string(ptr):
     if not ptr:
@@ -130,7 +144,11 @@ class StrArray:
     contents of 'struct' only remain valid within the StrArray context.
     """
 
-    def __init__(self, lst):
+    __array: 'GitStrrayC | ffi.NULL_TYPE'
+    __strings: list['None | ArrayC[char]']
+    __arr: 'ArrayC[char]'
+
+    def __init__(self, lst: None | Sequence[str | os.PathLike[str]]):
         # Allow passing in None as lg2 typically considers them the same as empty
         if lst is None:
             self.__array = ffi.NULL
@@ -139,29 +157,34 @@ def __init__(self, lst):
         if not isinstance(lst, (list, tuple)):
             raise TypeError('Value must be a list')
 
-        strings = [None] * len(lst)
+        strings: list[None | 'ArrayC[char]'] = [None] * len(lst)
         for i in range(len(lst)):
             li = lst[i]
             if not isinstance(li, str) and not hasattr(li, '__fspath__'):
                 raise TypeError('Value must be a string or PathLike object')
 
             strings[i] = ffi.new('char []', to_bytes(li))
 
-        self.__arr = ffi.new('char *[]', strings)
+        self.__arr = ffi.new('char *[]', strings)  # type: ignore[call-overload]
         self.__strings = strings
-        self.__array = ffi.new('git_strarray *', [self.__arr, len(strings)])
+        self.__array = ffi.new('git_strarray *', [self.__arr, len(strings)])  # type: ignore[call-overload]
 
-    def __enter__(self):
+    def __enter__(self) -> 'StrArray':
         return self
 
-    def __exit__(self, type, value, traceback):
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> None:
         pass
 
     @property
-    def ptr(self):
+    def ptr(self) -> 'GitStrrayC | ffi.NULL_TYPE':
         return self.__array
 
-    def assign_to(self, git_strarray):
+    def assign_to(self, git_strarray: 'GitStrrayC') -> None:
         if self.__array == ffi.NULL:
             git_strarray.strings = ffi.NULL
             git_strarray.count = 0