feat: adds typing

Now we have typing information added to the project and also
enabled checking in CI.

Author: kushaldas

.github/workflows/quality.yml

@@ -30,4 +30,10 @@ jobs:
         run: uv run ruff format --diff .
 
       - name: ruff check
-        run: uv run ruff check --diff .
+        run: uv run ruff check --diff .
+
+      - name: Install project dependencies for type checking
+        run: uv pip install asn1crypto
+
+      - name: ty type check
+        run: uv run ty check pkcs11/

MANIFEST.in

@@ -1,2 +1,3 @@
 graft extern/
-include pkcs11/*.pxd
+include pkcs11/*.pxd
+include pkcs11/py.typed

pkcs11/__init__.py

@@ -2,16 +2,24 @@
 :mod:`pkcs11` defines a high-level, "Pythonic" interface to PKCS#11.
 """
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
 from pkcs11.constants import *  # noqa: F403
 from pkcs11.exceptions import *  # noqa: F403
 from pkcs11.mechanisms import *  # noqa: F403
 from pkcs11.types import *  # noqa: F403
 from pkcs11.util import dh, dsa, ec, rsa, x509  # noqa: F401
 
-_loaded = {}
+if TYPE_CHECKING:
+    from pkcs11._pkcs11 import lib as _lib_type
+
+
+_loaded: dict[str, Any] = {}
 
 
-def lib(so):
+def lib(so: str) -> _lib_type:
     """
     Wrap the main library call coming from Cython with a preemptive
     dynamic loading.
@@ -34,7 +42,7 @@ def lib(so):
     return _lib
 
 
-def unload(so):
+def unload(so: str) -> None:
     global _loaded
     try:
         loaded_lib = _loaded[so]

pkcs11/_pkcs11.pyi

@@ -0,0 +1,44 @@
+"""Type stubs for the _pkcs11 Cython extension module."""
+
+from __future__ import annotations
+
+from typing import Any, Iterator
+
+from pkcs11.types import Slot, Token
+
+
+class lib:
+    """Main entry point for PKCS#11 library."""
+
+    so: str
+    manufacturer_id: str
+    library_description: str
+    initialized: bool
+
+    def __init__(self, so: str) -> None: ...
+
+    @property
+    def library_version(self) -> tuple[int, int]: ...
+
+    @property
+    def cryptoki_version(self) -> tuple[int, int]: ...
+
+    def initialize(self) -> None: ...
+    def finalize(self) -> None: ...
+    def reinitialize(self) -> None: ...
+    def unload(self) -> None: ...
+
+    def get_slots(self, token_present: bool = False) -> list[Slot]: ...
+    def get_tokens(
+        self,
+        token_label: str | None = None,
+        token_serial: bytes | None = None,
+        token_flags: Any | None = None,
+        slot_flags: Any | None = None,
+        mechanisms: Any | None = None,
+    ) -> Iterator[Token]: ...
+    def get_token(self, **kwargs: Any) -> Token: ...
+    def wait_for_slot_event(self, blocking: bool = True) -> Slot: ...
+
+    def __str__(self) -> str: ...
+    def __repr__(self) -> str: ...

pkcs11/attributes.py

@@ -1,36 +1,47 @@
+from __future__ import annotations
+
 from datetime import datetime
+from enum import IntEnum
 from struct import Struct
+from typing import Any, Callable, Final
 
-from pkcs11.constants import (
-    Attribute,
-    CertificateType,
-    MechanismFlag,
-    ObjectClass,
-)
+from pkcs11.constants import (Attribute, CertificateType, MechanismFlag,
+                              ObjectClass)
 from pkcs11.mechanisms import KeyType, Mechanism
 
+# Type aliases for pack/unpack function pairs
+PackFunc = Callable[[Any], bytes]
+UnpackFunc = Callable[[bytes], Any]
+Handler = tuple[PackFunc, UnpackFunc]
+
 # (Pack Function, Unpack Function) functions
-handle_bool = (Struct("?").pack, lambda v: False if len(v) == 0 else Struct("?").unpack(v)[0])
-handle_ulong = (Struct("L").pack, lambda v: Struct("L").unpack(v)[0])
-handle_str = (lambda s: s.encode("utf-8"), lambda b: b.decode("utf-8"))
-handle_date = (
+_bool_struct = Struct("?")
+_ulong_struct = Struct("L")
+
+handle_bool: Handler = (
+    _bool_struct.pack,
+    lambda v: False if len(v) == 0 else _bool_struct.unpack(v)[0],
+)
+handle_ulong: Handler = (_ulong_struct.pack, lambda v: _ulong_struct.unpack(v)[0])
+handle_str: Handler = (lambda s: s.encode("utf-8"), lambda b: b.decode("utf-8"))
+handle_date: Handler = (
     lambda s: s.strftime("%Y%m%d").encode("ascii"),
     lambda s: datetime.strptime(s.decode("ascii"), "%Y%m%d").date(),
 )
-handle_bytes = (bytes, bytes)
+handle_bytes: Handler = (bytes, bytes)
 # The PKCS#11 biginteger type is an array of bytes in network byte order.
 # If you have an int type, wrap it in biginteger()
-handle_biginteger = handle_bytes
+handle_biginteger: Handler = handle_bytes
 
 
-def _enum(type_):
+def _enum(type_: type[IntEnum]) -> Handler:
     """Factory to pack/unpack ints into IntEnums."""
     pack, unpack = handle_ulong
 
     return (lambda v: pack(int(v)), lambda v: type_(unpack(v)))
 
 
-ATTRIBUTE_TYPES = {
+ATTRIBUTE_TYPES: dict[Attribute, Handler] = {
     Attribute.ALWAYS_AUTHENTICATE: handle_bool,
     Attribute.ALWAYS_SENSITIVE: handle_bool,
     Attribute.APPLICATION: handle_str,
@@ -96,7 +107,7 @@ def _enum(type_):
 Map of attributes to (serialize, deserialize) functions.
 """
 
-ALL_CAPABILITIES = (
+ALL_CAPABILITIES: Final[tuple[Attribute, ...]] = (
     Attribute.ENCRYPT,
     Attribute.DECRYPT,
     Attribute.WRAP,
@@ -107,20 +118,29 @@ def _enum(type_):
 )
 
 
-def _apply_common(template, id_, label, store):
+def _apply_common(
+    template: dict[Attribute, Any],
+    id_: bytes | None,
+    label: str | None,
+    store: bool,
+) -> None:
     if id_:
         template[Attribute.ID] = id_
     if label:
         template[Attribute.LABEL] = label
     template[Attribute.TOKEN] = bool(store)
 
 
-def _apply_capabilities(template, possible_capas, capabilities):
+def _apply_capabilities(
+    template: dict[Attribute, Any],
+    possible_capas: tuple[Attribute, ...],
+    capabilities: MechanismFlag | int,
+) -> None:
     for attr in possible_capas:
         template[attr] = _capa_attr_to_mechanism_flag[attr] & capabilities
 
 
-_capa_attr_to_mechanism_flag = {
+_capa_attr_to_mechanism_flag: Final[dict[Attribute, MechanismFlag]] = {
     Attribute.ENCRYPT: MechanismFlag.ENCRYPT,
     Attribute.DECRYPT: MechanismFlag.DECRYPT,
     Attribute.WRAP: MechanismFlag.WRAP,
@@ -136,7 +156,12 @@ class AttributeMapper:
     Class mapping PKCS#11 attributes to and from Python values.
     """
 
-    def __init__(self):
+    attribute_types: dict[Attribute, Handler]
+    default_secret_key_template: dict[Attribute, Any]
+    default_public_key_template: dict[Attribute, Any]
+    default_private_key_template: dict[Attribute, Any]
+
+    def __init__(self) -> None:
         self.attribute_types = dict(ATTRIBUTE_TYPES)
         self.default_secret_key_template = {
             Attribute.CLASS: ObjectClass.SECRET_KEY,
@@ -158,33 +183,33 @@ def __init__(self):
             Attribute.SENSITIVE: True,
         }
 
-    def register_handler(self, key, pack, unpack):
+    def register_handler(self, key: Attribute, pack: PackFunc, unpack: UnpackFunc) -> None:
         self.attribute_types[key] = (pack, unpack)
 
-    def _handler(self, key):
+    def _handler(self, key: Attribute) -> Handler:
         try:
             return self.attribute_types[key]
         except KeyError as e:
             raise NotImplementedError(f"Can't handle attribute type {hex(key)}.") from e
 
-    def pack_attribute(self, key, value):
+    def pack_attribute(self, key: Attribute, value: Any) -> bytes:
         """Pack a Attribute value into a bytes array."""
         pack, _ = self._handler(key)
         return pack(value)
 
-    def unpack_attributes(self, key, value):
+    def unpack_attributes(self, key: Attribute, value: bytes) -> Any:
         """Unpack a Attribute bytes array into a Python value."""
         _, unpack = self._handler(key)
         return unpack(value)
 
     def public_key_template(
         self,
         *,
-        capabilities,
-        id_,
-        label,
-        store,
-    ):
+        capabilities: MechanismFlag | int,
+        id_: bytes | None,
+        label: str | None,
+        store: bool,
+    ) -> dict[Attribute, Any]:
         template = self.default_public_key_template
         _apply_capabilities(
             template, (Attribute.ENCRYPT, Attribute.WRAP, Attribute.VERIFY), capabilities
@@ -195,11 +220,11 @@ def public_key_template(
     def private_key_template(
         self,
         *,
-        capabilities,
-        id_,
-        label,
-        store,
-    ):
+        capabilities: MechanismFlag | int,
+        id_: bytes | None,
+        label: str | None,
+        store: bool,
+    ) -> dict[Attribute, Any]:
         template = self.default_private_key_template
         _apply_capabilities(
             template,
@@ -212,11 +237,11 @@ def private_key_template(
     def secret_key_template(
         self,
         *,
-        capabilities,
-        id_,
-        label,
-        store,
-    ):
+        capabilities: MechanismFlag | int,
+        id_: bytes | None,
+        label: str | None,
+        store: bool,
+    ) -> dict[Attribute, Any]:
         return self.generic_key_template(
             self.default_secret_key_template,
             capabilities=capabilities,
@@ -227,13 +252,13 @@ def secret_key_template(
 
     def generic_key_template(
         self,
-        base_template,
+        base_template: dict[Attribute, Any],
         *,
-        capabilities,
-        id_,
-        label,
-        store,
-    ):
+        capabilities: MechanismFlag | int,
+        id_: bytes | None,
+        label: str | None,
+        store: bool,
+    ) -> dict[Attribute, Any]:
         template = dict(base_template)
         _apply_capabilities(template, ALL_CAPABILITIES, capabilities)
         _apply_common(template, id_, label, store)

pkcs11/constants.py

@@ -5,9 +5,12 @@
 use these classes.
 """
 
+from __future__ import annotations
+
 from enum import IntEnum, IntFlag, unique
+from typing import Final
 
-DEFAULT = object()
+DEFAULT: Final[object] = object()
 """Sentinel value used in templates.
 
 Not all pkcs11 attribute sets are accepted by HSMs.
@@ -55,11 +58,11 @@ class ObjectClass(IntEnum):
 
     _VENDOR_DEFINED = 0x80000000
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return "<ObjectClass.%s>" % self.name
 
 
-_ARRAY_ATTRIBUTE = 0x40000000
+_ARRAY_ATTRIBUTE: Final[int] = 0x40000000
 """Attribute consists of an array of values."""
 
 
@@ -343,7 +346,7 @@ class Attribute(IntEnum):
 
     _VENDOR_DEFINED = 0x80000000
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return "<Attribute.%s>" % self.name