Return Python enums from _pygit2

Affected functions:
- Repository.status
- Repository.status_file
- Repository.merge_analysis
- DiffFile.flags
- DiffFile.mode
- DiffDelta.flags
- DiffDelta.status

Author: jorio

pygit2/__init__.py

@@ -59,6 +59,11 @@
 # libgit version tuple
 LIBGIT2_VER = (LIBGIT2_VER_MAJOR, LIBGIT2_VER_MINOR, LIBGIT2_VER_REVISION)
 
+# Let _pygit2 cache references to Python enum types.
+# This is separate from PyInit__pygit2() to avoid a circular import.
+cache_enums()
+del cache_enums  # Don't expose this to user code
+
 
 def init_repository(
         path: typing.Union[str, bytes, PathLike, None],

pygit2/_pygit2.pyi

@@ -1,18 +1,23 @@
 from typing import Iterator, Literal, Optional, overload
 from io import IOBase
 from . import Index, Submodule
-from .enums import ApplyLocation
-from .enums import BranchType
-from .enums import DiffFind
-from .enums import DiffOption
-from .enums import DiffStatsFormat
-from .enums import MergeAnalysis
-from .enums import MergePreference
-from .enums import ObjectType
-from .enums import Option
-from .enums import ReferenceFilter
-from .enums import ResetMode
-from .enums import SortMode
+from .enums import (
+    ApplyLocation,
+    BranchType,
+    DiffFind,
+    DiffFlag,
+    DiffOption,
+    DiffStatsFormat,
+    FileMode,
+    FileStatus,
+    MergeAnalysis,
+    MergePreference,
+    ObjectType,
+    Option,
+    ReferenceFilter,
+    ResetMode,
+    SortMode,
+)
 
 GIT_OBJ_BLOB: Literal[3]
 GIT_OBJ_COMMIT: Literal[1]
@@ -139,19 +144,19 @@ class Diff:
     def __len__(self) -> int: ...
 
 class DiffDelta:
-    flags: int
+    flags: DiffFlag
     is_binary: bool
-    new_file: DiffFile
     nfiles: int
+    new_file: DiffFile
     old_file: DiffFile
     similarity: int
-    status: int
+    status: FileStatus
     def status_char(self) -> str: ...
 
 class DiffFile:
-    flags: int
+    flags: DiffFlag
     id: Oid
-    mode: int
+    mode: FileMode
     path: str
     raw_path: bytes
     size: int

src/diff.c

@@ -47,6 +47,10 @@ extern PyTypeObject DiffLineType;
 extern PyTypeObject DiffStatsType;
 extern PyTypeObject RepositoryType;
 
+extern PyObject *DeltaStatusEnum;
+extern PyObject *DiffFlagEnum;
+extern PyObject *FileModeEnum;
+
 PyObject *
 wrap_diff(git_diff *diff, Repository *repo)
 {
@@ -204,13 +208,31 @@ DiffFile_from_c(DiffFile *dummy, PyObject *py_diff_file_ptr)
     return wrap_diff_file(diff_file);
 }
 
+PyDoc_STRVAR(DiffFile_flags__doc__,
+    "A combination of enums.DiffFlag constants."
+);
+
+PyObject *
+DiffFile_flags__get__(DiffFile *self)
+{
+    return pygit2_enum(DiffFlagEnum, self->flags);
+}
+
+PyDoc_STRVAR(DiffFile_mode__doc__,
+    "Mode of the entry (an enums.FileMode constant)."
+);
+
+PyObject *
+DiffFile_mode__get__(DiffFile *self)
+{
+    return pygit2_enum(FileModeEnum, self->mode);
+}
+
 PyMemberDef DiffFile_members[] = {
     MEMBER(DiffFile, id, T_OBJECT, "Oid of the item."),
     MEMBER(DiffFile, path, T_STRING, "Path to the entry."),
     MEMBER(DiffFile, raw_path, T_OBJECT, "Path to the entry (bytes)."),
     MEMBER(DiffFile, size, T_LONG, "Size of the entry."),
-    MEMBER(DiffFile, flags, T_UINT, "Combination of GIT_DIFF_FLAG_* flags."),
-    MEMBER(DiffFile, mode, T_USHORT, "Mode of the entry."),
     {NULL}
 };
 
@@ -219,6 +241,12 @@ PyMethodDef DiffFile_methods[] = {
     {NULL},
 };
 
+PyGetSetDef DiffFile_getsetters[] = {
+    GETTER(DiffFile, flags),
+    GETTER(DiffFile, mode),
+    {NULL},
+};
+
 PyDoc_STRVAR(DiffFile__doc__, "DiffFile object.");
 
 PyTypeObject DiffFileType = {
@@ -251,7 +279,7 @@ PyTypeObject DiffFileType = {
     0,                                         /* tp_iternext       */
     DiffFile_methods,                          /* tp_methods        */
     DiffFile_members,                          /* tp_members        */
-    0,                                         /* tp_getset         */
+    DiffFile_getsetters,                       /* tp_getset         */
     0,                                         /* tp_base           */
     0,                                         /* tp_dict           */
     0,                                         /* tp_descr_get      */
@@ -294,6 +322,26 @@ DiffDelta_is_binary__get__(DiffDelta *self)
     Py_RETURN_NONE;
 }
 
+PyDoc_STRVAR(DiffDelta_status__doc__,
+    "An enums.FileStatus constant."
+);
+
+PyObject *
+DiffDelta_status__get__(DiffDelta *self)
+{
+    return pygit2_enum(DeltaStatusEnum, self->status);
+}
+
+PyDoc_STRVAR(DiffDelta_flags__doc__,
+    "A combination of enums.DiffFlag constants."
+);
+
+PyObject *
+DiffDelta_flags__get__(DiffDelta *self)
+{
+    return pygit2_enum(DiffFlagEnum, self->flags);
+}
+
 static void
 DiffDelta_dealloc(DiffDelta *self)
 {
@@ -308,8 +356,6 @@ static PyMethodDef DiffDelta_methods[] = {
 };
 
 PyMemberDef DiffDelta_members[] = {
-    MEMBER(DiffDelta, status, T_UINT, "A GIT_DELTA_* constant."),
-    MEMBER(DiffDelta, flags, T_UINT, "Combination of GIT_DIFF_FLAG_* flags."),
     MEMBER(DiffDelta, similarity, T_USHORT, "For renamed and copied."),
     MEMBER(DiffDelta, nfiles, T_USHORT, "Number of files in the delta."),
     MEMBER(DiffDelta, old_file, T_OBJECT, "\"from\" side of the diff."),
@@ -319,6 +365,8 @@ PyMemberDef DiffDelta_members[] = {
 
 PyGetSetDef DiffDelta_getsetters[] = {
     GETTER(DiffDelta, is_binary),
+    GETTER(DiffDelta, status),
+    GETTER(DiffDelta, flags),
     {NULL}
 };
 

src/pygit2.c

@@ -42,6 +42,13 @@ PyObject *GitError;
 PyObject *AlreadyExistsError;
 PyObject *InvalidSpecError;
 
+PyObject *DeltaStatusEnum;
+PyObject *DiffFlagEnum;
+PyObject *FileModeEnum;
+PyObject *FileStatusEnum;
+PyObject *MergeAnalysisEnum;
+PyObject *MergePreferenceEnum;
+
 extern PyTypeObject RepositoryType;
 extern PyTypeObject OdbType;
 extern PyTypeObject OdbBackendType;
@@ -361,6 +368,74 @@ filter_unregister(PyObject *self, PyObject *args)
 }
 
 
+static void
+forget_enums(void)
+{
+    Py_XDECREF(DeltaStatusEnum);
+    Py_XDECREF(DiffFlagEnum);
+    Py_XDECREF(FileModeEnum);
+    Py_XDECREF(FileStatusEnum);
+    Py_XDECREF(MergeAnalysisEnum);
+    Py_XDECREF(MergePreferenceEnum);
+
+    DeltaStatusEnum = NULL;
+    DiffFlagEnum = NULL;
+    FileModeEnum = NULL;
+    FileStatusEnum = NULL;
+    MergeAnalysisEnum = NULL;
+    MergePreferenceEnum = NULL;
+}
+
+PyDoc_STRVAR(cache_enums__doc__,
+    "cache_enums()\n"
+    "\n"
+    "For internal use only. Do not call this from user code.\n"
+    "\n"
+    "Let the _pygit2 C module cache references to Python enums\n"
+    "defined in pygit2.enums.\n");
+
+PyObject *
+cache_enums(PyObject *self, PyObject *args)
+{
+    (void) args;
+
+    /* In case this is somehow being called several times, let go of old references */
+    forget_enums();
+
+    PyObject *enums = PyImport_ImportModule("pygit2.enums");
+    if (enums == NULL) {
+        return NULL;
+    }
+
+#define CACHE_PYGIT2_ENUM(name) do { \
+    name##Enum = PyObject_GetAttrString(enums, #name); \
+    if (name##Enum == NULL) { goto fail; } \
+} while (0)
+
+    CACHE_PYGIT2_ENUM(DeltaStatus);
+    CACHE_PYGIT2_ENUM(DiffFlag);
+    CACHE_PYGIT2_ENUM(FileMode);
+    CACHE_PYGIT2_ENUM(FileStatus);
+    CACHE_PYGIT2_ENUM(MergeAnalysis);
+    CACHE_PYGIT2_ENUM(MergePreference);
+
+#undef CACHE_PYGIT2_ENUM
+
+    Py_RETURN_NONE;
+
+fail:
+    Py_DECREF(enums);
+    forget_enums();
+    return NULL;
+}
+
+void
+free_module(void *self)
+{
+    forget_enums();
+}
+
+
 PyMethodDef module_methods[] = {
     {"discover_repository", discover_repository, METH_VARARGS, discover_repository__doc__},
     {"hash", hash, METH_VARARGS, hash__doc__},
@@ -371,10 +446,10 @@ PyMethodDef module_methods[] = {
     {"tree_entry_cmp", tree_entry_cmp, METH_VARARGS, tree_entry_cmp__doc__},
     {"filter_register", (PyCFunction)filter_register, METH_VARARGS | METH_KEYWORDS, filter_register__doc__},
     {"filter_unregister", filter_unregister, METH_VARARGS, filter_unregister__doc__},
+    {"cache_enums", cache_enums, METH_NOARGS, cache_enums__doc__},
     {NULL}
 };
 
-
 struct PyModuleDef moduledef = {
     PyModuleDef_HEAD_INIT,
     "_pygit2",                       /* m_name */
@@ -384,7 +459,7 @@ struct PyModuleDef moduledef = {
     NULL,                            /* m_reload */
     NULL,                            /* m_traverse */
     NULL,                            /* m_clear */
-    NULL,                            /* m_free */
+    free_module,                     /* m_free */
 };
 
 PyMODINIT_FUNC

src/repository.c

@@ -65,6 +65,10 @@ extern PyTypeObject NoteIterType;
 extern PyTypeObject StashType;
 extern PyTypeObject RefsIteratorType;
 
+extern PyObject *FileStatusEnum;
+extern PyObject *MergeAnalysisEnum;
+extern PyObject *MergePreferenceEnum;
+
 /* forward-declaration for Repsository._from_c() */
 PyTypeObject RepositoryType;
 
@@ -705,7 +709,22 @@ Repository_merge_analysis(Repository *self, PyObject *args)
         goto out;
     }
 
-    py_result = Py_BuildValue("(ii)", analysis, preference);
+    // Convert analysis to MergeAnalysis enum
+    PyObject *analysis_enum = pygit2_enum(MergeAnalysisEnum, analysis);
+    if (!analysis_enum) {
+        py_result = NULL;
+        goto out;
+    }
+
+    // Convert preference to MergePreference enum
+    PyObject *preference_enum = pygit2_enum(MergePreferenceEnum, preference);
+    if (!preference_enum) {
+        py_result = NULL;
+        Py_DECREF(analysis_enum);
+        goto out;
+    }
+
+    py_result = Py_BuildValue("(OO)", analysis_enum, preference_enum);
 
 out:
     git_reference_free(our_ref);
@@ -1754,14 +1773,17 @@ Repository_status(Repository *self, PyObject *args, PyObject *kw)
             path = entry->head_to_index->old_file.path;
         else
             path = entry->index_to_workdir->old_file.path;
-        status = PyLong_FromLong((long) entry->status);
+
+        /* Get corresponding entry in enums.FileStatus for status int */
+        status = pygit2_enum(FileStatusEnum, entry->status);
+        if (status == NULL)
+            goto error;
 
         err = PyDict_SetItemString(dict, path, status);
         Py_CLEAR(status);
 
         if (err < 0)
             goto error;
-
     }
 
     git_status_list_free(list);
@@ -1775,7 +1797,7 @@ Repository_status(Repository *self, PyObject *args, PyObject *kw)
 
 
 PyDoc_STRVAR(Repository_status_file__doc__,
-  "status_file(path: str) -> int\n"
+  "status_file(path: str) -> enums.FileStatus\n"
   "\n"
   "Returns the status of the given file path.");
 
@@ -1795,7 +1817,7 @@ Repository_status_file(Repository *self, PyObject *value)
     }
     free(path);
 
-    return PyLong_FromLong(status);
+    return pygit2_enum(FileStatusEnum, (int) status);
 }
 
 

src/utils.c

@@ -276,3 +276,18 @@ py_object_to_otype(PyObject *py_type)
     PyErr_SetString(PyExc_ValueError, "invalid target type");
     return GIT_OBJECT_INVALID; /* -1 */
 }
+
+
+/**
+ * Convert an integer to a reference to an IntEnum or IntFlag in pygit2.enums.
+ */
+PyObject *
+pygit2_enum(PyObject *enum_type, int value)
+{
+    if (!enum_type) {
+        PyErr_SetString(PyExc_TypeError, "an enum has not been cached in _pygit2.cache_enums()");
+        return NULL;
+    }
+    PyObject *enum_instance = PyObject_CallFunction(enum_type, "(i)", value);
+    return enum_instance;
+}

src/utils.h

@@ -99,6 +99,11 @@ char* pgit_encode_fsdefault(PyObject *value);
 
 git_otype py_object_to_otype(PyObject *py_type);
 
+
+/* Enum utilities (pygit2.enums) */
+PyObject *pygit2_enum(PyObject *enum_type, int value);
+
+
 /* Helpers to make shorter PyMethodDef and PyGetSetDef blocks */
 #define METHOD(type, name, args)\
   {#name, (PyCFunction) type ## _ ## name, args, type ## _ ## name ## __doc__}