Public API annotations and type checking with Pyrefly

.pre-commit-config.yaml

@@ -14,3 +14,9 @@ repos:
     hooks:
     -   id: ruff
         args: ["--fix", "--show-files"]
+-   repo: https://github.com/facebook/pyrefly-pre-commit
+    rev: 0.61.1
+    hooks:
+      - id: pyrefly-check
+        name: Pyrefly (type checking)
+        pass_filenames: false  # full repo checks

cloudpickle/cloudpickle.py

@@ -63,6 +63,8 @@
 import logging
 import opcode
 import pickle
+
+# pyrefly:ignore[missing-module-attribute]
 from pickle import _getattribute as _pickle_getattribute
 import platform
 import struct
@@ -79,6 +81,15 @@
 # cloudpickle. See: tests/test_backward_compat.py
 from types import CellType  # noqa: F401
 
+# always False; this is purely intended for type-checkers
+if typing.TYPE_CHECKING:
+    from _typeshed import SupportsWrite
+    from typing import Any, Callable
+    from typing_extensions import TypeAlias
+
+    # same as `_BufferCallback` in `typeshed/stdlib/_pickle.pyi`
+    _BufferCallback: TypeAlias = Callable[[pickle.PickleBuffer], Any] | None
+
 # cloudpickle is meant for inter process communication: we expect all
 # communicating processes to run the same Python version hence we favor
 # communication speed over compatibility:
@@ -124,7 +135,7 @@ def _lookup_class_or_track(class_tracker_id, class_def):
     return class_def
 
 
-def register_pickle_by_value(module):
+def register_pickle_by_value(module: types.ModuleType) -> None:
     """Register a module to make its functions and classes picklable by value.
 
     By default, functions and classes that are attributes of an importable
@@ -163,7 +174,7 @@ def register_pickle_by_value(module):
     _PICKLE_BY_VALUE_MODULES.add(module.__name__)
 
 
-def unregister_pickle_by_value(module):
+def unregister_pickle_by_value(module: types.ModuleType) -> None:
     """Unregister that the input module should be pickled by value."""
     if not isinstance(module, types.ModuleType):
         raise ValueError(f"Input should be a module object, got {str(module)} instead")
@@ -518,6 +529,7 @@ def _make_function(code, globals, name, argdefs, closure):
     return types.FunctionType(code, globals, name, argdefs, closure)
 
 
+@typing.no_type_check
 def _make_empty_cell():
     if False:
         # trick the compiler into creating an empty cell in our lambda
@@ -590,6 +602,7 @@ class id will also reuse this enum definition.
     return _lookup_class_or_track(class_tracker_id, enum_class)
 
 
+@typing.no_type_check
 def _make_typevar(name, bound, constraints, covariant, contravariant, class_tracker_id):
     tv = typing.TypeVar(
         name,
@@ -765,6 +778,7 @@ def _class_getstate(obj):
             # The abc caches and registered subclasses of a
             # class are bundled into the single _abc_impl attribute
             clsdict.pop("_abc_impl", None)
+            # pyrefly:ignore[missing-attribute]
             (registry, _, _, _) = abc._get_dump(obj)
 
             clsdict["_abc_impl"] = [subclass_weakref() for subclass_weakref in registry]
@@ -1220,8 +1234,11 @@ def _class_setstate(obj, state):
 
 
 _DATACLASSE_FIELD_TYPE_SENTINELS = {
+    # pyrefly:ignore[missing-attribute]
     dataclasses._FIELD.name: dataclasses._FIELD,
+    # pyrefly:ignore[missing-attribute]
     dataclasses._FIELD_CLASSVAR.name: dataclasses._FIELD_CLASSVAR,
+    # pyrefly:ignore[missing-attribute]
     dataclasses._FIELD_INITVAR.name: dataclasses._FIELD_INITVAR,
 }
 
@@ -1232,7 +1249,7 @@ def _get_dataclass_field_type_sentinel(name):
 
 class Pickler(pickle.Pickler):
     # set of reducers defined and used by cloudpickle (private)
-    _dispatch_table = {}
+    _dispatch_table: dict = {}
     _dispatch_table[classmethod] = _classmethod_reduce
     _dispatch_table[io.TextIOWrapper] = _file_reduce
     _dispatch_table[logging.Logger] = _logger_reduce
@@ -1258,8 +1275,10 @@ class Pickler(pickle.Pickler):
     _dispatch_table[abc.abstractclassmethod] = _classmethod_reduce
     _dispatch_table[abc.abstractstaticmethod] = _classmethod_reduce
     _dispatch_table[abc.abstractproperty] = _property_reduce
+    # pyrefly:ignore[missing-attribute]
     _dispatch_table[dataclasses._FIELD_BASE] = _dataclass_field_base_reduce
 
+    # pyrefly:ignore[bad-argument-type]
     dispatch_table = ChainMap(_dispatch_table, copyreg.dispatch_table)
 
     # function reducers are defined as instance methods of cloudpickle.Pickler
@@ -1317,14 +1336,19 @@ def _function_getnewargs(self, func):
 
         return code, base_globals, None, None, closure
 
-    def dump(self, obj):
+    def dump(self, obj: object) -> None:
         try:
             return super().dump(obj)
         except RecursionError as e:
             msg = "Could not pickle object as excessively deep recursion required."
             raise pickle.PicklingError(msg) from e
 
-    def __init__(self, file, protocol=None, buffer_callback=None):
+    def __init__(
+        self,
+        file: "SupportsWrite[bytes]",
+        protocol: "int | None" = None,
+        buffer_callback: "_BufferCallback" = None,
+    ) -> None:
         if protocol is None:
             protocol = DEFAULT_PROTOCOL
         super().__init__(file, protocol=protocol, buffer_callback=buffer_callback)
@@ -1412,8 +1436,9 @@ def reducer_override(self, obj):
         # Pickler's types.FunctionType and type savers. Note: the type saver
         # must override Pickler.save_global, because pickle.py contains a
         # hard-coded call to save_global when pickling meta-classes.
-        dispatch = pickle.Pickler.dispatch.copy()
+        dispatch = pickle.Pickler.dispatch.copy()  # pyrefly:ignore[missing-attribute]
 
+        @typing.no_type_check
         def _save_reduce_pickle5(
             self,
             func,
@@ -1447,6 +1472,7 @@ def _save_reduce_pickle5(
             # the stack.
             write(pickle.POP)
 
+        @typing.no_type_check
         def save_global(self, obj, name=None, pack=struct.pack):
             """Main dispatch method.
 
@@ -1473,6 +1499,7 @@ def save_global(self, obj, name=None, pack=struct.pack):
 
         dispatch[type] = save_global
 
+        @typing.no_type_check
         def save_function(self, obj, name=None):
             """Registered with the dispatch to handle all function types.
 
@@ -1488,6 +1515,7 @@ def save_function(self, obj, name=None):
                     *self._dynamic_function_reduce(obj), obj=obj
                 )
 
+        @typing.no_type_check
         def save_pypy_builtin_func(self, obj):
             """Save pypy equivalent of builtin functions.
 
@@ -1519,7 +1547,12 @@ def save_pypy_builtin_func(self, obj):
 # Shorthands similar to pickle.dump/pickle.dumps
 
 
-def dump(obj, file, protocol=None, buffer_callback=None):
+def dump(
+    obj: object,
+    file: "SupportsWrite[bytes]",
+    protocol: "int | None" = None,
+    buffer_callback: "_BufferCallback" = None,
+) -> None:
     """Serialize obj as bytes streamed into file
 
     protocol defaults to cloudpickle.DEFAULT_PROTOCOL which is an alias to
@@ -1535,7 +1568,11 @@ def dump(obj, file, protocol=None, buffer_callback=None):
     Pickler(file, protocol=protocol, buffer_callback=buffer_callback).dump(obj)
 
 
-def dumps(obj, protocol=None, buffer_callback=None):
+def dumps(
+    obj: object,
+    protocol: "int | None" = None,
+    buffer_callback: "_BufferCallback" = None,
+) -> bytes:
     """Serialize obj as a string of bytes allocated in memory
 
     protocol defaults to cloudpickle.DEFAULT_PROTOCOL which is an alias to

dev-requirements.txt

@@ -19,3 +19,5 @@ coverage
 ./tests/cloudpickle_testpkg
 # Required for setup of the above utility package:
 setuptools; python_version >= '3.12'
+# type-checking
+pyrefly==0.61.1

pyproject.toml

@@ -39,3 +39,22 @@ preview = true
 [tool.ruff]
 line-length = 88
 target-version = "py38"
+
+[tool.pyrefly]
+project-includes = ["cloudpickle"]
+enabled-ignores = ["pyrefly"]
+
+[tool.pyrefly.errors]
+implicit-abstract-class = "error"
+implicitly-defined-attribute = "error"
+not-required-key-access = "error"
+open-unpacking = "error"
+unannotated-attribute = "error"
+untyped-import = "error"
+unused-ignore = "error"
+variance-mismatch = "error"
+# TODO: enable these once everything is annotated
+missing-override-decorator = "ignore"
+implicit-any = "ignore"
+unannotated-parameter = "ignore"
+unannotated-return = "ignore"