Skip to content

Commit 6ba20a6

Browse files
kzndotshkzndotsh
committed
feat(typings): add type stubs for py-pglite modules
- Introduced type stubs for the py-pglite library, including core components such as clients, configuration, manager, and SQLAlchemy integration. - Implemented abstract database client interfaces and specific implementations for psycopg and asyncpg. - Added configuration dataclass for PGlite settings and utility functions for database operations. - Enhanced SQLAlchemy integration with specific fixtures and manager functionalities for improved testing support. - Established a comprehensive typing structure to facilitate better development and testing practices.
1 parent b00e5eb commit 6ba20a6

File tree

10 files changed

+668
-0
lines changed

10 files changed

+668
-0
lines changed

‎typings/py_pglite/__init__.pyi‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
"""
2+
This type stub file was generated by pyright.
3+
"""
4+
5+
from .clients import AsyncpgClient, PsycopgClient, get_client, get_default_client
6+
from .config import PGliteConfig
7+
from .manager import PGliteManager
8+
9+
"""py-pglite: Python testing library for PGlite integration.
10+
11+
Provides seamless integration between PGlite (in-memory PostgreSQL)
12+
and Python test suites with support for SQLAlchemy, SQLModel, and Django.
13+
"""
14+
__version__ = ...
15+
__all__ = ["PGliteConfig", "PGliteManager", "get_client", "get_default_client", "PsycopgClient", "AsyncpgClient"]

‎typings/py_pglite/clients.pyi‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
"""
2+
This type stub file was generated by pyright.
3+
"""
4+
5+
from abc import ABC, abstractmethod
6+
from typing import Any
7+
8+
"""Database client abstraction for py-pglite.
9+
10+
Provides unified interface for both psycopg and asyncpg clients,
11+
allowing users to choose their preferred PostgreSQL driver.
12+
"""
13+
logger = ...
14+
class DatabaseClient(ABC):
15+
"""Abstract database client interface."""
16+
@abstractmethod
17+
def connect(self, connection_string: str) -> Any:
18+
"""Create a connection to the database."""
19+
...
20+
21+
@abstractmethod
22+
def execute_query(self, connection: Any, query: str, params: Any = ...) -> list[tuple]:
23+
"""Execute a query and return results."""
24+
...
25+
26+
@abstractmethod
27+
def test_connection(self, connection_string: str) -> bool:
28+
"""Test if database connection is working."""
29+
...
30+
31+
@abstractmethod
32+
def get_database_version(self, connection_string: str) -> str | None:
33+
"""Get PostgreSQL version string."""
34+
...
35+
36+
@abstractmethod
37+
def close_connection(self, connection: Any) -> None:
38+
"""Close a database connection."""
39+
...
40+
41+
42+
43+
class PsycopgClient(DatabaseClient):
44+
"""psycopg-based database client."""
45+
def __init__(self) -> None:
46+
...
47+
48+
def connect(self, connection_string: str) -> Any:
49+
"""Create a psycopg connection."""
50+
...
51+
52+
def execute_query(self, connection: Any, query: str, params: Any = ...) -> list[tuple]:
53+
"""Execute query using psycopg."""
54+
...
55+
56+
def test_connection(self, connection_string: str) -> bool:
57+
"""Test psycopg connection."""
58+
...
59+
60+
def get_database_version(self, connection_string: str) -> str | None:
61+
"""Get PostgreSQL version using psycopg."""
62+
...
63+
64+
def close_connection(self, connection: Any) -> None:
65+
"""Close psycopg connection."""
66+
...
67+
68+
69+
70+
class AsyncpgClient(DatabaseClient):
71+
"""asyncpg-based database client."""
72+
def __init__(self) -> None:
73+
...
74+
75+
def connect(self, connection_string: str) -> Any:
76+
"""Create an asyncpg connection (sync wrapper)."""
77+
...
78+
79+
def execute_query(self, connection: Any, query: str, params: Any = ...) -> list[tuple]:
80+
"""Execute query using asyncpg (sync wrapper)."""
81+
...
82+
83+
def test_connection(self, connection_string: str) -> bool:
84+
"""Test asyncpg connection."""
85+
...
86+
87+
def get_database_version(self, connection_string: str) -> str | None:
88+
"""Get PostgreSQL version using asyncpg."""
89+
...
90+
91+
def close_connection(self, connection: Any) -> None:
92+
"""Close asyncpg connection."""
93+
...
94+
95+
96+
97+
def get_default_client() -> DatabaseClient:
98+
"""Get the default database client.
99+
100+
Prefers psycopg if available, falls back to asyncpg.
101+
"""
102+
...
103+
104+
def get_client(client_type: str = ...) -> DatabaseClient:
105+
"""Get a database client by type.
106+
107+
Args:
108+
client_type: "psycopg", "asyncpg", or "auto" (default)
109+
110+
Returns:
111+
DatabaseClient instance
112+
"""
113+
...
114+
115+
__all__ = ["DatabaseClient", "PsycopgClient", "AsyncpgClient", "get_default_client", "get_client"]

‎typings/py_pglite/config.pyi‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
"""
2+
This type stub file was generated by pyright.
3+
"""
4+
5+
from dataclasses import dataclass
6+
from pathlib import Path
7+
8+
"""Configuration for PGlite testing."""
9+
@dataclass
10+
class PGliteConfig:
11+
"""Configuration for PGlite test database.
12+
13+
Args:
14+
timeout: Timeout in seconds for PGlite startup (default: 30)
15+
cleanup_on_exit: Whether to cleanup socket/process on exit (default: True)
16+
log_level: Logging level for PGlite operations (default: "INFO")
17+
socket_path: Custom socket path (default: secure temp directory)
18+
work_dir: Working directory for PGlite files (default: None, uses temp)
19+
node_modules_check: Whether to verify node_modules exists (default: True)
20+
auto_install_deps: Whether to auto-install npm dependencies (default: True)
21+
extensions: List of PGlite extensions to enable (e.g., ["pgvector"])
22+
node_options: Custom NODE_OPTIONS for the Node.js process
23+
"""
24+
timeout: int = ...
25+
cleanup_on_exit: bool = ...
26+
log_level: str = ...
27+
socket_path: str = ...
28+
work_dir: Path | None = ...
29+
node_modules_check: bool = ...
30+
auto_install_deps: bool = ...
31+
extensions: list[str] | None = ...
32+
node_options: str | None = ...
33+
def __post_init__(self) -> None:
34+
"""Validate configuration after initialization."""
35+
...
36+
37+
@property
38+
def log_level_int(self) -> int:
39+
"""Get logging level as integer."""
40+
...
41+
42+
def get_connection_string(self) -> str:
43+
"""Get PostgreSQL connection string for SQLAlchemy usage."""
44+
...
45+
46+
def get_psycopg_uri(self) -> str:
47+
"""Get PostgreSQL URI for direct psycopg usage."""
48+
...
49+
50+
def get_dsn(self) -> str:
51+
"""Get PostgreSQL DSN connection string for direct psycopg usage."""
52+
...
53+
54+
55+

‎typings/py_pglite/extensions.pyi‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
"""
2+
This type stub file was generated by pyright.
3+
"""
4+
5+
"""Extension management for py-pglite.
6+
7+
This module provides a registry of supported PGlite extensions and the
8+
necessary JavaScript import details for each.
9+
"""
10+
SUPPORTED_EXTENSIONS: dict[str, dict[str, str]] = ...

‎typings/py_pglite/manager.pyi‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
"""
2+
This type stub file was generated by pyright.
3+
"""
4+
5+
from typing import Any
6+
from .config import PGliteConfig
7+
8+
"""Core PGlite process management."""
9+
class PGliteManager:
10+
"""Manages PGlite process lifecycle for testing.
11+
12+
Framework-agnostic PGlite process manager. Provides database connections
13+
through framework-specific methods that require their respective dependencies.
14+
"""
15+
def __init__(self, config: PGliteConfig | None = ...) -> None:
16+
"""Initialize PGlite manager.
17+
18+
Args:
19+
config: Configuration for PGlite. If None, uses defaults.
20+
"""
21+
...
22+
23+
def __enter__(self) -> PGliteManager:
24+
"""Context manager entry."""
25+
...
26+
27+
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
28+
"""Context manager exit."""
29+
...
30+
31+
def start(self) -> None:
32+
"""Start the PGlite server."""
33+
...
34+
35+
def stop(self) -> None:
36+
"""Stop the PGlite server."""
37+
...
38+
39+
def is_running(self) -> bool:
40+
"""Check if PGlite process is running."""
41+
...
42+
43+
def get_connection_string(self) -> str:
44+
"""Get the database connection string for framework-agnostic usage.
45+
46+
Returns:
47+
PostgreSQL connection string
48+
49+
Raises:
50+
RuntimeError: If PGlite server is not running
51+
"""
52+
...
53+
54+
def get_dsn(self) -> str:
55+
"""Get the database DSN string for framework-agnostic usage.
56+
57+
Returns:
58+
PostgreSQL DSN string
59+
"""
60+
...
61+
62+
def wait_for_ready_basic(self, max_retries: int = ..., delay: float = ...) -> bool:
63+
"""Wait for database to be ready using framework-agnostic connection test.
64+
65+
Args:
66+
max_retries: Maximum number of connection attempts
67+
delay: Delay between attempts in seconds
68+
69+
Returns:
70+
True if database becomes ready, False otherwise
71+
"""
72+
...
73+
74+
def wait_for_ready(self, max_retries: int = ..., delay: float = ...) -> bool:
75+
"""Wait for database to be ready (framework-agnostic).
76+
77+
This is an alias for wait_for_ready_basic() to maintain API consistency
78+
across different manager types while keeping the base manager framework-agnostic.
79+
80+
Args:
81+
max_retries: Maximum number of connection attempts
82+
delay: Delay between attempts in seconds
83+
84+
Returns:
85+
True if database becomes ready, False otherwise
86+
"""
87+
...
88+
89+
def restart(self) -> None:
90+
"""Restart the PGlite server.
91+
92+
Stops the current server if running and starts a new one.
93+
"""
94+
...
95+
96+
def get_psycopg_uri(self) -> str:
97+
"""Get the database URI for psycopg usage.
98+
99+
Returns:
100+
PostgreSQL URI string compatible with psycopg
101+
102+
Raises:
103+
RuntimeError: If PGlite server is not running
104+
"""
105+
...
106+
107+
108+

‎typings/py_pglite/sqlalchemy/__init__.pyi‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
"""
2+
This type stub file was generated by pyright.
3+
"""
4+
5+
from .fixtures import pglite_engine, pglite_session, pglite_sqlalchemy_engine, pglite_sqlalchemy_session
6+
from .manager import SQLAlchemyPGliteManager
7+
from .utils import create_all_tables, drop_all_tables, get_session_class
8+
9+
"""SQLAlchemy integration for py-pglite.
10+
11+
This module provides SQLAlchemy-specific fixtures and utilities for py-pglite.
12+
"""
13+
__all__ = ["SQLAlchemyPGliteManager", "pglite_engine", "pglite_session", "pglite_sqlalchemy_session", "pglite_sqlalchemy_engine", "create_all_tables", "drop_all_tables", "get_session_class"]

‎typings/py_pglite/sqlalchemy/fixtures.pyi‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
"""
2+
This type stub file was generated by pyright.
3+
"""
4+
5+
import pytest
6+
from collections.abc import Generator
7+
from typing import Any
8+
from sqlalchemy.engine import Engine
9+
from sqlalchemy.orm import Session
10+
from ..config import PGliteConfig
11+
from .manager import SQLAlchemyPGliteManager
12+
13+
"""SQLAlchemy-specific pytest fixtures for PGlite integration."""
14+
HAS_SQLMODEL = ...
15+
logger = ...
16+
@pytest.fixture(scope="session")
17+
def pglite_config() -> PGliteConfig:
18+
"""Pytest fixture providing PGlite configuration."""
19+
...
20+
21+
@pytest.fixture(scope="session")
22+
def pglite_sqlalchemy_manager(pglite_config: PGliteConfig) -> Generator[SQLAlchemyPGliteManager, None, None]:
23+
"""Pytest fixture providing an SQLAlchemy-enabled PGlite manager."""
24+
...
25+
26+
@pytest.fixture(scope="session")
27+
def pglite_engine(pglite_sqlalchemy_manager: SQLAlchemyPGliteManager) -> Engine:
28+
"""Pytest fixture providing a SQLAlchemy engine connected to PGlite.
29+
30+
Uses the SQLAlchemy-enabled manager to ensure proper SQLAlchemy integration.
31+
"""
32+
...
33+
34+
@pytest.fixture(scope="session")
35+
def pglite_sqlalchemy_engine(pglite_sqlalchemy_manager: SQLAlchemyPGliteManager) -> Engine:
36+
"""Pytest fixture providing an optimized SQLAlchemy engine connected to PGlite."""
37+
...
38+
39+
@pytest.fixture(scope="function")
40+
def pglite_session(pglite_engine: Engine) -> Generator[Any, None, None]:
41+
"""Pytest fixture providing a SQLAlchemy/SQLModel session with proper isolation.
42+
43+
This fixture ensures database isolation between tests by cleaning all data
44+
at the start of each test.
45+
"""
46+
...
47+
48+
@pytest.fixture(scope="function")
49+
def pglite_sqlalchemy_session(pglite_session: Session) -> Session:
50+
"""Legacy fixture name for backwards compatibility."""
51+
...
52+

0 commit comments

Comments
 (0)