PyO3: Add Python stub file for typing support (#587)

* add stubs file for rust library type hints

* return python dictionaries directly instead of pyobject

* add docstrings

* fill stubs file

* add integration test for stack data

* format stubs file

* ignore upgrade lints in pyi file so it is py 3.9 compatible

* fix return type for vault init

* docstring tweaks

.github/workflows/integration.yml

@@ -339,6 +339,12 @@ jobs:
       - name: Check Python vault package
         run: python -m pip show nitor-vault
 
+      - name: Check stack status with Python library
+        run: python -c "from n_vault import Vault; print(Vault().stack_status())"
+
+      - name: Check stack status with Rust CLI
+        run: bin/rust/vault stack
+
       - name: Delete all keys with Python library
         run: python -c "from n_vault import Vault; Vault().delete_many(Vault().list_all())"
 

python-pyo3/pyproject.toml

@@ -98,6 +98,10 @@ exclude = [
     "venv*",
 ]
 
+[tool.ruff.lint.per-file-ignores]
+# Remove after upgrading minimun Python version to 3.10
+"nitor_vault_rs.pyi" = ["UP"]
+
 [tool.ruff.lint.isort]
 # https://docs.astral.sh/ruff/settings/#isort
 combine-as-imports = true

python-pyo3/python/n_vault/nitor_vault_rs.pyi

@@ -0,0 +1,171 @@
+from typing import Any, Dict, List, Optional
+
+class VaultConfig:
+    """
+    Optional parameters for a `Vault` instance.
+
+    Attributes:
+        vault_stack (Optional[str]): The name of the CloudFormation stack.
+        region (Optional[str]): The AWS region for the bucket.
+        bucket (Optional[str]): The name of the S3 bucket.
+        key (Optional[str]): The encryption key ARN.
+        prefix (Optional[str]): The prefix for keys.
+        profile (Optional[str]): The AWS profile name.
+        iam_id (Optional[str]): The IAM user ID.
+        iam_secret (Optional[str]): The IAM secret key.
+    """
+
+    vault_stack: Optional[str]
+    region: Optional[str]
+    bucket: Optional[str]
+    key: Optional[str]
+    prefix: Optional[str]
+    profile: Optional[str]
+    iam_id: Optional[str]
+    iam_secret: Optional[str]
+
+    def __init__(
+        self,
+        vault_stack: Optional[str] = None,
+        region: Optional[str] = None,
+        bucket: Optional[str] = None,
+        key: Optional[str] = None,
+        prefix: Optional[str] = None,
+        profile: Optional[str] = None,
+        iam_id: Optional[str] = None,
+        iam_secret: Optional[str] = None,
+    ) -> VaultConfig:
+        """
+        Initialize a VaultConfig instance with optional parameters.
+
+        Args:
+            vault_stack: The name of the CloudFormation stack.
+            region: The AWS region for the bucket.
+            bucket: The name of the S3 bucket.
+            key: The encryption key ARN.
+            prefix: The prefix for keys.
+            profile: The AWS profile name.
+            iam_id: The IAM user ID.
+            iam_secret: The IAM secret key.
+        """
+        ...
+
+def delete(name: str, config: VaultConfig) -> None:
+    """
+    Delete data in S3 for the given key name.
+    """
+
+def delete_many(names: List[str], config: VaultConfig) -> None:
+    """
+    Delete data for multiple keys.
+    """
+
+def direct_decrypt(data: bytes, config: VaultConfig) -> bytes:
+    """
+    Decrypt data with KMS.
+
+    Args:
+        data: Encrypted bytes to decrypt.
+        config: Vault configuration.
+
+    Returns:
+        Decrypted bytes.
+    """
+
+def direct_encrypt(data: bytes, config: VaultConfig) -> bytes:
+    """
+    Encrypt data with KMS.
+
+    Args:
+        data: Plaintext bytes to encrypt.
+        config: Vault configuration.
+
+    Returns:
+        Encrypted bytes.
+    """
+
+def exists(name: str, config: VaultConfig) -> bool:
+    """
+    Check if the given key name exists in the S3 bucket.
+
+    Args:
+        name: The key name to check.
+        config: Vault configuration.
+
+    Returns:
+        True if the key exists, False otherwise.
+    """
+
+def init(config: VaultConfig) -> Dict[str, Any]:
+    """
+    Initialize a new Vault stack.
+
+    Args:
+        config: Vault configuration.
+
+    Returns:
+        A dictionary containing stack initialization details.
+    """
+
+def list_all(config: VaultConfig) -> List[str]:
+    """
+    Get all available secrets.
+
+    Args:
+        config: Vault configuration.
+
+    Returns:
+        A list of key names.
+    """
+
+def lookup(name: str, config: VaultConfig) -> bytes:
+    """
+    Lookup the value for the given key name.
+
+    Args:
+        name: The key name to look up.
+        config: Vault configuration.
+
+    Returns:
+        The raw bytes stored under the given key.
+    """
+
+def run(args: List[str]) -> None:
+    """
+    Run Vault CLI with the given arguments.
+
+    Args:
+        args: List of command-line arguments, including program name.
+    """
+
+def stack_status(config: VaultConfig) -> Dict[str, Any]:
+    """
+    Get the Vault CloudFormation stack status.
+
+    Args:
+        config: Vault configuration.
+
+    Returns:
+        A dictionary with the stack status details.
+    """
+
+def store(name: str, value: bytes, config: VaultConfig) -> None:
+    """
+    Store an encrypted value with the given key name in S3.
+
+    Args:
+        name: Key name for the data.
+        value: Bytes to store.
+        config: Vault configuration.
+    """
+
+def update(config: VaultConfig) -> Dict[str, Any]:
+    """
+    Update the Vault CloudFormation stack with the current template.
+
+    Args:
+        config: Vault configuration.
+
+    Returns:
+        A dictionary with the stack update details.
+    """

python-pyo3/python/n_vault/vault.py

@@ -59,7 +59,7 @@ class Vault:
 
     Note that initializing this class only saves the optional parameters,
     but does *not* construct an actual vault instance.
-    Each method in this class creates its own Vault instance internally in the Rust library.
+    Each method in this class creates its own Vault instance internally in the Rust library if needed.
     """
 
     def __init__(

python-pyo3/src/lib.rs

@@ -13,6 +13,7 @@ use nitor_vault::{CreateStackResult, UpdateStackResult, Value, Vault};
 static RUNTIME: LazyLock<Runtime> =
     LazyLock::new(|| Runtime::new().expect("Failed to start async runtime."));
 
+/// Optional parameters for a `Vault` instance.
 #[pyclass]
 #[derive(Debug, Default, Clone)]
 pub struct VaultConfig {
@@ -77,6 +78,7 @@ impl From<VaultConfig> for RustVaultConfig {
     }
 }
 
+/// Delete data in S3 for given key name.
 #[pyfunction()]
 fn delete(name: &str, config: VaultConfig) -> PyResult<()> {
     RUNTIME.block_on(async {
@@ -89,6 +91,7 @@ fn delete(name: &str, config: VaultConfig) -> PyResult<()> {
     })
 }
 
+/// Delete data for multiple keys.
 #[pyfunction()]
 #[allow(clippy::needless_pass_by_value)]
 fn delete_many(names: Vec<String>, config: VaultConfig) -> PyResult<()> {
@@ -102,6 +105,7 @@ fn delete_many(names: Vec<String>, config: VaultConfig) -> PyResult<()> {
     })
 }
 
+/// Decrypt data with KMS.
 #[pyfunction()]
 fn direct_decrypt(data: &[u8], config: VaultConfig) -> PyResult<Cow<[u8]>> {
     // Returns Cow<[u8]> instead of Vec since that will get mapped to bytes for the Python side
@@ -118,6 +122,7 @@ fn direct_decrypt(data: &[u8], config: VaultConfig) -> PyResult<Cow<[u8]>> {
     })
 }
 
+/// Encrypt data with KMS.
 #[pyfunction()]
 fn direct_encrypt(data: &[u8], config: VaultConfig) -> PyResult<Cow<[u8]>> {
     RUNTIME.block_on(async {
@@ -132,6 +137,9 @@ fn direct_encrypt(data: &[u8], config: VaultConfig) -> PyResult<Cow<[u8]>> {
     })
 }
 
+/// Check if the given key name already exists in the S3 bucket.
+///
+/// Returns True if the key exists, False otherwise.
 #[pyfunction()]
 fn exists(name: &str, config: VaultConfig) -> PyResult<bool> {
     RUNTIME.block_on(async {
@@ -146,8 +154,9 @@ fn exists(name: &str, config: VaultConfig) -> PyResult<bool> {
     })
 }
 
+/// Initialize new Vault stack.
 #[pyfunction()]
-fn init(config: VaultConfig) -> PyResult<PyObject> {
+fn init(config: VaultConfig) -> PyResult<Py<PyDict>> {
     let result = RUNTIME.block_on(async {
         Vault::init(
             config.vault_stack,
@@ -184,6 +193,9 @@ fn init(config: VaultConfig) -> PyResult<PyObject> {
     })
 }
 
+/// Get all available secrets.
+///
+/// Returns a list of key names.
 #[pyfunction()]
 fn list_all(config: VaultConfig) -> PyResult<Vec<String>> {
     RUNTIME.block_on(async {
@@ -198,6 +210,9 @@ fn list_all(config: VaultConfig) -> PyResult<Vec<String>> {
     })
 }
 
+/// Lookup value for given key name.
+///
+/// Returns raw bytes.
 #[pyfunction()]
 fn lookup(name: &str, config: VaultConfig) -> PyResult<Cow<[u8]>> {
     RUNTIME.block_on(async {
@@ -214,17 +229,18 @@ fn lookup(name: &str, config: VaultConfig) -> PyResult<Cow<[u8]>> {
     })
 }
 
-#[pyfunction]
 /// Run Vault CLI with given args.
+#[pyfunction]
 fn run(args: Vec<String>) -> PyResult<()> {
     RUNTIME.block_on(async {
         nitor_vault::run_cli_with_args(args).await?;
         Ok(())
     })
 }
 
+/// Get vault Cloudformation stack status.
 #[pyfunction()]
-fn stack_status(config: VaultConfig) -> PyResult<PyObject> {
+fn stack_status(config: VaultConfig) -> PyResult<Py<PyDict>> {
     let data = RUNTIME.block_on(async {
         Vault::from_config(config.into())
             .await
@@ -240,6 +256,7 @@ fn stack_status(config: VaultConfig) -> PyResult<PyObject> {
     })
 }
 
+/// Store encrypted value with given key name in S3.
 #[pyfunction()]
 fn store(name: &str, value: &[u8], config: VaultConfig) -> PyResult<()> {
     RUNTIME.block_on(async {
@@ -254,8 +271,9 @@ fn store(name: &str, value: &[u8], config: VaultConfig) -> PyResult<()> {
     })
 }
 
+/// Update the vault Cloudformation stack with the current template.
 #[pyfunction()]
-fn update(config: VaultConfig) -> PyResult<PyObject> {
+fn update(config: VaultConfig) -> PyResult<Py<PyDict>> {
     let result = RUNTIME.block_on(async {
         Vault::from_config(config.into())
             .await