Add comprehensive type stubs to improve Python import

.gitignore

@@ -160,3 +160,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 .idea/
+
+# VSCode
+.vscode/

pyproject.toml

@@ -38,3 +38,4 @@ tests = [
 [tool.maturin]
 name = "vodozemac"
 bindings = "pyo3"
+python-source = "."

src/lib.rs

@@ -10,6 +10,17 @@ use error::*;
 use pyo3::{prelude::*, types::PyBytes};
 
 #[pymodule(name = "vodozemac")]
+/// Python bindings for the vodozemac Rust library.
+///
+/// This library provides Python bindings for vodozemac, a pure Rust
+/// implementation of the Matrix cryptographic protocols including:
+/// - Olm (end-to-end encryption for 1:1 conversations)
+/// - Megolm (end-to-end encryption for group conversations)
+/// - SAS (Short Authentication String) verification
+/// - Public key encryption (PK encryption) for key backup
+///
+/// All the classes and functions in this module are thread-safe and can be used
+/// in concurrent environments.
 fn my_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<account::Account>()?;
     m.add_class::<session::Session>()?;

vodozemac/__init__.py

@@ -0,0 +1,44 @@
+import contextlib
+from .vodozemac import *  # type: ignore
+
+# Define all available exports for better IDE support
+__all__ = [
+    # Classes
+    "Account",
+    "Session", 
+    "AnyOlmMessage",
+    "PreKeyMessage",
+    "Sas",
+    "EstablishedSas", 
+    "GroupSession",
+    "InboundGroupSession",
+    "SessionKey",
+    "ExportedSessionKey", 
+    "MegolmMessage",
+    "Ed25519PublicKey",
+    "Ed25519Signature",
+    "Curve25519PublicKey", 
+    "Curve25519SecretKey",
+    "PkDecryption",
+    "PkEncryption",
+    "Message",
+    # Exceptions
+    "KeyException",
+    "SignatureException", 
+    "DecodeException",
+    "LibolmPickleException",
+    "SessionKeyDecodeException",
+    "PickleException",
+    "SessionCreationException",
+    "SasException", 
+    "OlmDecryptionException",
+    "MegolmDecryptionException",
+    "PkInvalidKeySizeException",
+    "PkDecodeException",
+]
+
+with contextlib.suppress(ImportError):
+    from . import vodozemac  # type: ignore
+    __doc__ = vodozemac.__doc__
+    if hasattr(vodozemac, "__all__"):
+        __all__ = vodozemac.__all__ 

vodozemac/__init__.pyi

@@ -0,0 +1,252 @@
+"""Type stubs for vodozemac - Python bindings for the vodozemac Rust library."""
+
+from typing import Optional, Dict, Any, Tuple, List
+from typing_extensions import Self
+
+__all__ = [
+    "Account", "Session", "AnyOlmMessage", "PreKeyMessage", "Sas", "EstablishedSas", 
+    "GroupSession", "InboundGroupSession", "SessionKey", "ExportedSessionKey", 
+    "MegolmMessage", "Ed25519PublicKey", "Ed25519Signature", "Curve25519PublicKey", 
+    "Curve25519SecretKey", "PkDecryption", "PkEncryption", "Message", "KeyException",
+    "SignatureException", "DecodeException", "LibolmPickleException", "SessionKeyDecodeException",
+    "PickleException", "SessionCreationException", "SasException", "OlmDecryptionException", 
+    "MegolmDecryptionException", "PkInvalidKeySizeException", "PkDecodeException"
+]
+
+# Exceptions
+class KeyException(ValueError): ...
+class SignatureException(ValueError): ...
+class DecodeException(ValueError): ...
+class LibolmPickleException(ValueError): ...
+class SessionKeyDecodeException(ValueError): ...
+class PickleException(ValueError): ...
+class SessionCreationException(ValueError): ...
+class SasException(ValueError): ...
+class OlmDecryptionException(ValueError): ...
+class MegolmDecryptionException(ValueError): ...
+class PkInvalidKeySizeException(ValueError): ...
+class PkDecodeException(ValueError): ...
+
+# Key Types
+class Ed25519PublicKey:
+    """An Ed25519 public key."""
+
+    @classmethod
+    def from_base64(cls, key: str) -> Ed25519PublicKey: ...
+    def to_base64(self) -> str: ...
+    def verify_signature(self, message: bytes, signature: Ed25519Signature) -> None: ...
+    def __eq__(self, other: object) -> bool: ...
+
+class Ed25519Signature:
+    """An Ed25519 signature."""
+
+    @classmethod
+    def from_base64(cls, signature: str) -> Ed25519Signature: ...
+    def to_base64(self) -> str: ...
+
+class Curve25519PublicKey:
+    """A Curve25519 public key."""
+
+    @classmethod
+    def from_base64(cls, key: str) -> Curve25519PublicKey: ...
+    @classmethod
+    def from_bytes(cls, bytes: bytes) -> Curve25519PublicKey: ...
+    def to_base64(self) -> str: ...
+    def to_bytes(self) -> bytes: ...
+    def __eq__(self, other: object) -> bool: ...
+
+class Curve25519SecretKey:
+    """A Curve25519 secret key."""
+
+    def __init__(self) -> None: ...
+    @classmethod
+    def from_base64(cls, key: str) -> Curve25519SecretKey: ...
+    @classmethod
+    def from_bytes(cls, bytes: bytes) -> Curve25519SecretKey: ...
+    def to_base64(self) -> str: ...
+    def to_bytes(self) -> bytes: ...
+    def public_key(self) -> Curve25519PublicKey: ...
+
+# Session Keys
+class SessionKey:
+    """A Megolm session key."""
+
+    def __init__(self, session_key: str) -> None: ...
+    def to_base64(self) -> str: ...
+
+class ExportedSessionKey:
+    """An exported Megolm session key."""
+
+    def __init__(self, session_key: str) -> None: ...
+    def to_base64(self) -> str: ...
+
+# Messages
+class AnyOlmMessage:
+    """A message encrypted using the Olm protocol."""
+
+    @classmethod
+    def pre_key(cls, message: bytes) -> AnyOlmMessage: ...
+    @classmethod
+    def normal(cls, message: bytes) -> AnyOlmMessage: ...
+    @classmethod
+    def from_parts(cls, message_type: int, ciphertext: bytes) -> AnyOlmMessage: ...
+    def to_pre_key(self) -> Optional[PreKeyMessage]: ...
+    def to_parts(self) -> Tuple[int, bytes]: ...
+
+class PreKeyMessage:
+    """A pre-key message for the Olm protocol."""
+
+    @classmethod
+    def from_base64(cls, message: str) -> PreKeyMessage: ...
+    def to_any(self) -> AnyOlmMessage: ...
+    def session_id(self) -> str: ...
+
+class MegolmMessage:
+    """A Megolm encrypted message."""
+
+    @classmethod
+    def from_base64(cls, message: str) -> MegolmMessage: ...
+    @classmethod
+    def from_bytes(cls, message: bytes) -> MegolmMessage: ...
+    def to_base64(self) -> str: ...
+    def to_bytes(self) -> bytes: ...
+    def message_index(self) -> int: ...
+    def signature(self) -> Ed25519Signature: ...
+
+# Account and Session
+class Account:
+    """An Olm account."""
+
+    def __init__(self) -> None: ...
+    @classmethod
+    def from_pickle(cls, pickle: str, pickle_key: bytes) -> Account: ...
+    @classmethod
+    def from_libolm_pickle(cls, pickle: str, pickle_key: bytes) -> Account: ...
+    def pickle(self, pickle_key: bytes) -> str: ...
+
+    @property
+    def ed25519_key(self) -> Ed25519PublicKey: ...
+    @property
+    def curve25519_key(self) -> Curve25519PublicKey: ...
+    @property
+    def one_time_keys(self) -> Dict[str, Curve25519PublicKey]: ...
+    @property
+    def max_number_of_one_time_keys(self) -> int: ...
+
+    def sign(self, message: bytes) -> Ed25519Signature: ...
+    def generate_one_time_keys(self, count: int) -> None: ...
+    def mark_keys_as_published(self) -> None: ...
+    def create_outbound_session(self, identity_key: Curve25519PublicKey, one_time_key: Curve25519PublicKey) -> Session: ...
+    def create_inbound_session(self, message: PreKeyMessage) -> Session: ...
+
+class Session:
+    """An Olm session."""
+
+    @classmethod
+    def from_pickle(cls, pickle: str, pickle_key: bytes) -> Session: ...
+    @classmethod
+    def from_libolm_pickle(cls, pickle: str, pickle_key: bytes) -> Session: ...
+
+    @property
+    def session_id(self) -> str: ...
+
+    def pickle(self, pickle_key: bytes) -> str: ...
+    def session_matches(self, message: PreKeyMessage) -> bool: ...
+    def encrypt(self, plaintext: bytes) -> AnyOlmMessage: ...
+    def decrypt(self, message: AnyOlmMessage) -> bytes: ...
+
+# Group Sessions
+class GroupSession:
+    """An outbound Megolm group session."""
+
+    def __init__(self) -> None: ...
+    @classmethod
+    def from_pickle(cls, pickle: str, pickle_key: bytes) -> GroupSession: ...
+
+    @property
+    def session_id(self) -> str: ...
+    @property
+    def message_index(self) -> int: ...
+    @property
+    def session_key(self) -> SessionKey: ...
+
+    def encrypt(self, plaintext: bytes) -> MegolmMessage: ...
+    def pickle(self, pickle_key: bytes) -> str: ...
+
+class DecryptedMessage:
+    """A decrypted Megolm message."""
+    plaintext: bytes
+    message_index: int
+
+class InboundGroupSession:
+    """An inbound Megolm group session."""
+
+    def __init__(self, session_key: SessionKey) -> None: ...
+    @classmethod
+    def import_session(cls, session_key: ExportedSessionKey) -> InboundGroupSession: ...
+    @classmethod
+    def from_pickle(cls, pickle: str, pickle_key: bytes) -> InboundGroupSession: ...
+
+    @property
+    def session_id(self) -> str: ...
+    @property
+    def first_known_index(self) -> int: ...
+
+    def export_at(self, index: int) -> Optional[ExportedSessionKey]: ...
+    def decrypt(self, message: MegolmMessage) -> DecryptedMessage: ...
+    def pickle(self, pickle_key: bytes) -> str: ...
+
+# SAS (Short Authentication String)
+class Sas:
+    """A SAS object for interactive key verification."""
+
+    def __init__(self) -> None: ...
+    @property
+    def public_key(self) -> Curve25519PublicKey: ...
+    def diffie_hellman(self, key: Curve25519PublicKey) -> EstablishedSas: ...
+
+class EstablishedSas:
+    """An established SAS session."""
+
+    def bytes(self, info: str) -> SasBytes: ...
+    def calculate_mac_invalid_base64(self, input: str, info: str) -> str: ...
+    def calculate_mac(self, input: str, info: str) -> str: ...
+    def verify_mac(self, input: str, info: str, tag: str) -> None: ...
+
+class SasBytes:
+    """SAS bytes for generating verification codes."""
+
+    @property
+    def emoji_indices(self) -> List[int]: ...  # Actually returns [u8; 7] but typing as List[int]
+    @property
+    def decimals(self) -> Tuple[int, int, int]: ...
+
+# PK Encryption (Public Key Encryption)
+class Message:
+    """A message encrypted using PK encryption."""
+
+    def __init__(self, ciphertext: bytes, mac: bytes, ephemeral_key: bytes) -> None: ...
+    @classmethod
+    def from_base64(cls, ciphertext: str, mac: str, ephemeral_key: str) -> Message: ...
+    def to_base64(self) -> Tuple[str, str, str]: ...
+
+    ciphertext: bytes
+    mac: bytes
+    ephemeral_key: bytes
+
+class PkDecryption:
+    """PK decryption object."""
+
+    def __init__(self) -> None: ...
+    @classmethod
+    def from_key(cls, key: Curve25519SecretKey) -> PkDecryption: ...
+    @property
+    def public_key(self) -> Curve25519PublicKey: ...
+    def decrypt(self, message: Message) -> bytes: ...
+
+class PkEncryption:
+    """PK encryption object."""
+
+    @classmethod
+    def from_key(cls, key: Curve25519PublicKey) -> PkEncryption: ...
+    def encrypt(self, message: bytes) -> Message: ... 

vodozemac/py.typed

@@ -0,0 +1 @@
+