From b8a8651c232180d15233cfb679cc996793eb8cc0 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 19 Sep 2024 08:40:31 -0700 Subject: [PATCH] ci: Enable mypy checks for checkpoint-postgres lib --- libs/checkpoint-postgres/Makefile | 3 +- .../langgraph/checkpoint/postgres/__init__.py | 22 +++++--- .../langgraph/checkpoint/postgres/aio.py | 24 +++++--- .../langgraph/checkpoint/postgres/base.py | 18 +++--- libs/checkpoint-postgres/poetry.lock | 56 +++++++++---------- libs/checkpoint-postgres/pyproject.toml | 10 ++++ libs/checkpoint-postgres/tests/conftest.py | 8 ++- libs/checkpoint-postgres/tests/test_async.py | 16 +++--- libs/checkpoint-postgres/tests/test_sync.py | 16 +++--- 9 files changed, 103 insertions(+), 70 deletions(-) diff --git a/libs/checkpoint-postgres/Makefile b/libs/checkpoint-postgres/Makefile index 7057f1a97..aadb22cd7 100644 --- a/libs/checkpoint-postgres/Makefile +++ b/libs/checkpoint-postgres/Makefile @@ -41,7 +41,8 @@ lint lint_diff lint_package lint_tests: poetry run ruff check . [ "$(PYTHON_FILES)" = "" ] || poetry run ruff format $(PYTHON_FILES) --diff [ "$(PYTHON_FILES)" = "" ] || poetry run ruff check --select I $(PYTHON_FILES) - [ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE) || poetry run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE) + [ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE) + [ "$(PYTHON_FILES)" = "" ] || poetry run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE) format format_diff: poetry run ruff format $(PYTHON_FILES) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 6d5fbac8c..821b74af6 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -5,7 +5,7 @@ from typing import Any, Iterator, List, Optional, Union from langchain_core.runnables import RunnableConfig from psycopg import Connection, Cursor, Pipeline from psycopg.errors import UndefinedTable -from psycopg.rows import dict_row +from psycopg.rows import DictRow, dict_row from psycopg.types.json import Jsonb from psycopg_pool import ConnectionPool @@ -22,9 +22,11 @@ from langgraph.checkpoint.postgres.base import ( ) from langgraph.checkpoint.serde.base import SerializerProtocol +Conn = Union[Connection[DictRow], ConnectionPool[Connection[DictRow]]] + @contextmanager -def _get_connection(conn: Union[Connection, ConnectionPool]) -> Iterator[Connection]: +def _get_connection(conn: Conn) -> Iterator[Connection[DictRow]]: if isinstance(conn, Connection): yield conn elif isinstance(conn, ConnectionPool): @@ -39,7 +41,7 @@ class PostgresSaver(BasePostgresSaver): def __init__( self, - conn: Union[Connection, ConnectionPool], + conn: Conn, pipe: Optional[Pipeline] = None, serde: Optional[SerializerProtocol] = None, ) -> None: @@ -85,9 +87,13 @@ class PostgresSaver(BasePostgresSaver): """ with self._cursor() as cur: try: - version = cur.execute( + row = cur.execute( "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1" - ).fetchone()["v"] + ).fetchone() + if row is None: + version = -1 + else: + version = row["v"] except UndefinedTable: version = -1 for v, migration in zip( @@ -212,7 +218,7 @@ class PostgresSaver(BasePostgresSaver): checkpoint_id = get_checkpoint_id(config) checkpoint_ns = config["configurable"].get("checkpoint_ns", "") if checkpoint_id: - args = (thread_id, checkpoint_ns, checkpoint_id) + args: tuple[Any, ...] = (thread_id, checkpoint_ns, checkpoint_id) where = "WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s" else: args = (thread_id, checkpoint_ns) @@ -306,7 +312,7 @@ class PostgresSaver(BasePostgresSaver): self._dump_blobs( thread_id, checkpoint_ns, - copy.pop("channel_values"), + copy.pop("channel_values"), # type: ignore[misc] new_versions, ), ) @@ -356,7 +362,7 @@ class PostgresSaver(BasePostgresSaver): ) @contextmanager - def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor]: + def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]: with _get_connection(self.conn) as conn: if self.pipe: # a connection in pipeline mode can be used concurrently diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 9aa08ad0f..43a8cfe47 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -5,7 +5,7 @@ from typing import Any, AsyncIterator, Iterator, List, Optional, Union from langchain_core.runnables import RunnableConfig from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline from psycopg.errors import UndefinedTable -from psycopg.rows import dict_row +from psycopg.rows import DictRow, dict_row from psycopg.types.json import Jsonb from psycopg_pool import AsyncConnectionPool @@ -20,11 +20,13 @@ from langgraph.checkpoint.base import ( from langgraph.checkpoint.postgres.base import BasePostgresSaver from langgraph.checkpoint.serde.base import SerializerProtocol +Conn = Union[AsyncConnection[DictRow], AsyncConnectionPool[AsyncConnection[DictRow]]] + @asynccontextmanager async def _get_connection( - conn: Union[AsyncConnection, AsyncConnectionPool], -) -> AsyncIterator[AsyncConnection]: + conn: Conn, +) -> AsyncIterator[AsyncConnection[DictRow]]: if isinstance(conn, AsyncConnection): yield conn elif isinstance(conn, AsyncConnectionPool): @@ -39,7 +41,7 @@ class AsyncPostgresSaver(BasePostgresSaver): def __init__( self, - conn: Union[AsyncConnection, AsyncConnectionPool], + conn: Conn, pipe: Optional[AsyncPipeline] = None, serde: Optional[SerializerProtocol] = None, ) -> None: @@ -93,7 +95,11 @@ class AsyncPostgresSaver(BasePostgresSaver): results = await cur.execute( "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1" ) - version = (await results.fetchone())["v"] + row = await results.fetchone() + if row is None: + version = -1 + else: + version = row["v"] except UndefinedTable: version = -1 for v, migration in zip( @@ -180,7 +186,7 @@ class AsyncPostgresSaver(BasePostgresSaver): checkpoint_id = get_checkpoint_id(config) checkpoint_ns = config["configurable"].get("checkpoint_ns", "") if checkpoint_id: - args = (thread_id, checkpoint_ns, checkpoint_id) + args: tuple[Any, ...] = (thread_id, checkpoint_ns, checkpoint_id) where = "WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s" else: args = (thread_id, checkpoint_ns) @@ -265,7 +271,7 @@ class AsyncPostgresSaver(BasePostgresSaver): self._dump_blobs, thread_id, checkpoint_ns, - copy.pop("channel_values"), + copy.pop("channel_values"), # type: ignore[misc] new_versions, ), ) @@ -314,7 +320,9 @@ class AsyncPostgresSaver(BasePostgresSaver): await cur.executemany(query, params) @asynccontextmanager - async def _cursor(self, *, pipeline: bool = False) -> AsyncIterator[AsyncCursor]: + async def _cursor( + self, *, pipeline: bool = False + ) -> AsyncIterator[AsyncCursor[DictRow]]: async with _get_connection(self.conn) as conn: if self.pipe: # a connection in pipeline mode can be used concurrently diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 282eb2c86..f2fa21eba 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -1,5 +1,5 @@ import random -from typing import Any, List, Optional, Tuple +from typing import Any, List, Optional, Tuple, cast from langchain_core.runnables import RunnableConfig from psycopg.types.json import Jsonb @@ -7,7 +7,9 @@ from psycopg.types.json import Jsonb from langgraph.checkpoint.base import ( WRITES_IDX_MAP, BaseCheckpointSaver, + ChannelVersions, Checkpoint, + CheckpointMetadata, get_checkpoint_id, ) from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer @@ -122,7 +124,7 @@ INSERT_CHECKPOINT_WRITES_SQL = """ """ -class BasePostgresSaver(BaseCheckpointSaver): +class BasePostgresSaver(BaseCheckpointSaver[str]): SELECT_SQL = SELECT_SQL MIGRATIONS = MIGRATIONS UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL @@ -165,8 +167,8 @@ class BasePostgresSaver(BaseCheckpointSaver): thread_id: str, checkpoint_ns: str, values: dict[str, Any], - versions: dict[str, str], - ) -> list[tuple[str, str, str, str, str, bytes]]: + versions: ChannelVersions, + ) -> list[tuple[str, str, str, str, str, Optional[bytes]]]: if not versions: return [] @@ -175,7 +177,7 @@ class BasePostgresSaver(BaseCheckpointSaver): thread_id, checkpoint_ns, k, - ver, + cast(str, ver), *( self.serde.dumps_typed(values[k]) if k in values @@ -208,7 +210,7 @@ class BasePostgresSaver(BaseCheckpointSaver): checkpoint_id: str, task_id: str, writes: list[tuple[str, Any]], - ) -> list[tuple[str, str, str, int, str, str, bytes]]: + ) -> list[tuple[str, str, str, str, int, str, str, bytes]]: return [ ( thread_id, @@ -222,10 +224,10 @@ class BasePostgresSaver(BaseCheckpointSaver): for idx, (channel, value) in enumerate(writes) ] - def _load_metadata(self, metadata: dict[str, Any]) -> dict[str, Any]: + def _load_metadata(self, metadata: dict[str, Any]) -> CheckpointMetadata: return self.jsonplus_serde.loads(self.jsonplus_serde.dumps(metadata)) - def _dump_metadata(self, metadata) -> str: + def _dump_metadata(self, metadata: CheckpointMetadata) -> str: serialized_metadata = self.jsonplus_serde.dumps(metadata) return serialized_metadata.decode() diff --git a/libs/checkpoint-postgres/poetry.lock b/libs/checkpoint-postgres/poetry.lock index 4116e8257..be01d9a74 100644 --- a/libs/checkpoint-postgres/poetry.lock +++ b/libs/checkpoint-postgres/poetry.lock @@ -434,38 +434,38 @@ files = [ [[package]] name = "mypy" -version = "1.11.0" +version = "1.11.2" description = "Optional static typing for Python" optional = false python-versions = ">=3.8" files = [ - {file = "mypy-1.11.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:a3824187c99b893f90c845bab405a585d1ced4ff55421fdf5c84cb7710995229"}, - {file = "mypy-1.11.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:96f8dbc2c85046c81bcddc246232d500ad729cb720da4e20fce3b542cab91287"}, - {file = "mypy-1.11.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1a5d8d8dd8613a3e2be3eae829ee891b6b2de6302f24766ff06cb2875f5be9c6"}, - {file = "mypy-1.11.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:72596a79bbfb195fd41405cffa18210af3811beb91ff946dbcb7368240eed6be"}, - {file = "mypy-1.11.0-cp310-cp310-win_amd64.whl", hash = "sha256:35ce88b8ed3a759634cb4eb646d002c4cef0a38f20565ee82b5023558eb90c00"}, - {file = "mypy-1.11.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:98790025861cb2c3db8c2f5ad10fc8c336ed2a55f4daf1b8b3f877826b6ff2eb"}, - {file = "mypy-1.11.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:25bcfa75b9b5a5f8d67147a54ea97ed63a653995a82798221cca2a315c0238c1"}, - {file = "mypy-1.11.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0bea2a0e71c2a375c9fa0ede3d98324214d67b3cbbfcbd55ac8f750f85a414e3"}, - {file = "mypy-1.11.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:d2b3d36baac48e40e3064d2901f2fbd2a2d6880ec6ce6358825c85031d7c0d4d"}, - {file = "mypy-1.11.0-cp311-cp311-win_amd64.whl", hash = "sha256:d8e2e43977f0e09f149ea69fd0556623919f816764e26d74da0c8a7b48f3e18a"}, - {file = "mypy-1.11.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:1d44c1e44a8be986b54b09f15f2c1a66368eb43861b4e82573026e04c48a9e20"}, - {file = "mypy-1.11.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:cea3d0fb69637944dd321f41bc896e11d0fb0b0aa531d887a6da70f6e7473aba"}, - {file = "mypy-1.11.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a83ec98ae12d51c252be61521aa5731f5512231d0b738b4cb2498344f0b840cd"}, - {file = "mypy-1.11.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:c7b73a856522417beb78e0fb6d33ef89474e7a622db2653bc1285af36e2e3e3d"}, - {file = "mypy-1.11.0-cp312-cp312-win_amd64.whl", hash = "sha256:f2268d9fcd9686b61ab64f077be7ffbc6fbcdfb4103e5dd0cc5eaab53a8886c2"}, - {file = "mypy-1.11.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:940bfff7283c267ae6522ef926a7887305945f716a7704d3344d6d07f02df850"}, - {file = "mypy-1.11.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:14f9294528b5f5cf96c721f231c9f5b2733164e02c1c018ed1a0eff8a18005ac"}, - {file = "mypy-1.11.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d7b54c27783991399046837df5c7c9d325d921394757d09dbcbf96aee4649fe9"}, - {file = "mypy-1.11.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:65f190a6349dec29c8d1a1cd4aa71284177aee5949e0502e6379b42873eddbe7"}, - {file = "mypy-1.11.0-cp38-cp38-win_amd64.whl", hash = "sha256:dbe286303241fea8c2ea5466f6e0e6a046a135a7e7609167b07fd4e7baf151bf"}, - {file = "mypy-1.11.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:104e9c1620c2675420abd1f6c44bab7dd33cc85aea751c985006e83dcd001095"}, - {file = "mypy-1.11.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:f006e955718ecd8d159cee9932b64fba8f86ee6f7728ca3ac66c3a54b0062abe"}, - {file = "mypy-1.11.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:becc9111ca572b04e7e77131bc708480cc88a911adf3d0239f974c034b78085c"}, - {file = "mypy-1.11.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:6801319fe76c3f3a3833f2b5af7bd2c17bb93c00026a2a1b924e6762f5b19e13"}, - {file = "mypy-1.11.0-cp39-cp39-win_amd64.whl", hash = "sha256:c1a184c64521dc549324ec6ef7cbaa6b351912be9cb5edb803c2808a0d7e85ac"}, - {file = "mypy-1.11.0-py3-none-any.whl", hash = "sha256:56913ec8c7638b0091ef4da6fcc9136896914a9d60d54670a75880c3e5b99ace"}, - {file = "mypy-1.11.0.tar.gz", hash = "sha256:93743608c7348772fdc717af4aeee1997293a1ad04bc0ea6efa15bf65385c538"}, + {file = "mypy-1.11.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:d42a6dd818ffce7be66cce644f1dff482f1d97c53ca70908dff0b9ddc120b77a"}, + {file = "mypy-1.11.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:801780c56d1cdb896eacd5619a83e427ce436d86a3bdf9112527f24a66618fef"}, + {file = "mypy-1.11.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:41ea707d036a5307ac674ea172875f40c9d55c5394f888b168033177fce47383"}, + {file = "mypy-1.11.2-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:6e658bd2d20565ea86da7d91331b0eed6d2eee22dc031579e6297f3e12c758c8"}, + {file = "mypy-1.11.2-cp310-cp310-win_amd64.whl", hash = "sha256:478db5f5036817fe45adb7332d927daa62417159d49783041338921dcf646fc7"}, + {file = "mypy-1.11.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:75746e06d5fa1e91bfd5432448d00d34593b52e7e91a187d981d08d1f33d4385"}, + {file = "mypy-1.11.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:a976775ab2256aadc6add633d44f100a2517d2388906ec4f13231fafbb0eccca"}, + {file = "mypy-1.11.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:cd953f221ac1379050a8a646585a29574488974f79d8082cedef62744f0a0104"}, + {file = "mypy-1.11.2-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:57555a7715c0a34421013144a33d280e73c08df70f3a18a552938587ce9274f4"}, + {file = "mypy-1.11.2-cp311-cp311-win_amd64.whl", hash = "sha256:36383a4fcbad95f2657642a07ba22ff797de26277158f1cc7bd234821468b1b6"}, + {file = "mypy-1.11.2-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:e8960dbbbf36906c5c0b7f4fbf2f0c7ffb20f4898e6a879fcf56a41a08b0d318"}, + {file = "mypy-1.11.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:06d26c277962f3fb50e13044674aa10553981ae514288cb7d0a738f495550b36"}, + {file = "mypy-1.11.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6e7184632d89d677973a14d00ae4d03214c8bc301ceefcdaf5c474866814c987"}, + {file = "mypy-1.11.2-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:3a66169b92452f72117e2da3a576087025449018afc2d8e9bfe5ffab865709ca"}, + {file = "mypy-1.11.2-cp312-cp312-win_amd64.whl", hash = "sha256:969ea3ef09617aff826885a22ece0ddef69d95852cdad2f60c8bb06bf1f71f70"}, + {file = "mypy-1.11.2-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:37c7fa6121c1cdfcaac97ce3d3b5588e847aa79b580c1e922bb5d5d2902df19b"}, + {file = "mypy-1.11.2-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:4a8a53bc3ffbd161b5b2a4fff2f0f1e23a33b0168f1c0778ec70e1a3d66deb86"}, + {file = "mypy-1.11.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2ff93107f01968ed834f4256bc1fc4475e2fecf6c661260066a985b52741ddce"}, + {file = "mypy-1.11.2-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:edb91dded4df17eae4537668b23f0ff6baf3707683734b6a818d5b9d0c0c31a1"}, + {file = "mypy-1.11.2-cp38-cp38-win_amd64.whl", hash = "sha256:ee23de8530d99b6db0573c4ef4bd8f39a2a6f9b60655bf7a1357e585a3486f2b"}, + {file = "mypy-1.11.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:801ca29f43d5acce85f8e999b1e431fb479cb02d0e11deb7d2abb56bdaf24fd6"}, + {file = "mypy-1.11.2-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:af8d155170fcf87a2afb55b35dc1a0ac21df4431e7d96717621962e4b9192e70"}, + {file = "mypy-1.11.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f7821776e5c4286b6a13138cc935e2e9b6fde05e081bdebf5cdb2bb97c9df81d"}, + {file = "mypy-1.11.2-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:539c570477a96a4e6fb718b8d5c3e0c0eba1f485df13f86d2970c91f0673148d"}, + {file = "mypy-1.11.2-cp39-cp39-win_amd64.whl", hash = "sha256:3f14cd3d386ac4d05c5a39a51b84387403dadbd936e17cb35882134d4f8f0d24"}, + {file = "mypy-1.11.2-py3-none-any.whl", hash = "sha256:b499bc07dbdcd3de92b0a8b29fdf592c111276f6a12fe29c30f6c417dd546d12"}, + {file = "mypy-1.11.2.tar.gz", hash = "sha256:7f9993ad3e0ffdc95c2a14b66dee63729f021968bff8ad911867579c65d13a79"}, ] [package.dependencies] diff --git a/libs/checkpoint-postgres/pyproject.toml b/libs/checkpoint-postgres/pyproject.toml index 605720b91..59592acd4 100644 --- a/libs/checkpoint-postgres/pyproject.toml +++ b/libs/checkpoint-postgres/pyproject.toml @@ -51,3 +51,13 @@ lint.select = [ "I", # isort ] lint.ignore = ["E501", "B008", "UP007", "UP006"] + +[tool.mypy] +# https://mypy.readthedocs.io/en/stable/config_file.html +disallow_untyped_defs = "True" +explicit_package_bases = "True" +warn_no_return = "False" +warn_unused_ignores = "True" +warn_redundant_casts = "True" +allow_redefinition = "True" +disable_error_code = "typeddict-item, return-value" diff --git a/libs/checkpoint-postgres/tests/conftest.py b/libs/checkpoint-postgres/tests/conftest.py index 6100be1d4..36336643f 100644 --- a/libs/checkpoint-postgres/tests/conftest.py +++ b/libs/checkpoint-postgres/tests/conftest.py @@ -1,13 +1,15 @@ +from typing import AsyncIterator + import pytest from psycopg import AsyncConnection from psycopg.errors import UndefinedTable -from psycopg.rows import dict_row +from psycopg.rows import DictRow, dict_row DEFAULT_URI = "postgres://postgres:postgres@localhost:5441/postgres?sslmode=disable" @pytest.fixture(scope="function") -async def conn(): +async def conn() -> AsyncIterator[AsyncConnection[DictRow]]: async with await AsyncConnection.connect( DEFAULT_URI, autocommit=True, prepare_threshold=0, row_factory=dict_row ) as conn: @@ -15,7 +17,7 @@ async def conn(): @pytest.fixture(scope="function", autouse=True) -async def clear_test_db(conn): +async def clear_test_db(conn: AsyncConnection[DictRow]) -> None: """Delete all tables before each test.""" try: await conn.execute("DELETE FROM checkpoints") diff --git a/libs/checkpoint-postgres/tests/test_async.py b/libs/checkpoint-postgres/tests/test_async.py index 6f9f7d78b..288996529 100644 --- a/libs/checkpoint-postgres/tests/test_async.py +++ b/libs/checkpoint-postgres/tests/test_async.py @@ -1,5 +1,7 @@ +from typing import Any + import pytest -from conftest import DEFAULT_URI +from conftest import DEFAULT_URI # type: ignore from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( @@ -13,7 +15,7 @@ from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver class TestAsyncPostgresSaver: @pytest.fixture(autouse=True) - async def setup(self): + async def setup(self) -> None: # objects for test setup self.config_1: RunnableConfig = { "configurable": { @@ -58,20 +60,20 @@ class TestAsyncPostgresSaver: async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver: await saver.setup() - async def test_asearch(self): + async def test_asearch(self) -> None: async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver: await saver.aput(self.config_1, self.chkpnt_1, self.metadata_1, {}) await saver.aput(self.config_2, self.chkpnt_2, self.metadata_2, {}) await saver.aput(self.config_3, self.chkpnt_3, self.metadata_3, {}) # call method / assertions - query_1: CheckpointMetadata = {"source": "input"} # search by 1 key - query_2: CheckpointMetadata = { + query_1 = {"source": "input"} # search by 1 key + query_2 = { "step": 1, "writes": {"foo": "bar"}, } # search by multiple keys - query_3: CheckpointMetadata = {} # search by no keys, return all checkpoints - query_4: CheckpointMetadata = {"source": "update", "step": 1} # no match + query_3: dict[str, Any] = {} # search by no keys, return all checkpoints + query_4 = {"source": "update", "step": 1} # no match search_results_1 = [c async for c in saver.alist(None, filter=query_1)] assert len(search_results_1) == 1 diff --git a/libs/checkpoint-postgres/tests/test_sync.py b/libs/checkpoint-postgres/tests/test_sync.py index a2fbcbd88..f64059a28 100644 --- a/libs/checkpoint-postgres/tests/test_sync.py +++ b/libs/checkpoint-postgres/tests/test_sync.py @@ -1,5 +1,7 @@ +from typing import Any + import pytest -from conftest import DEFAULT_URI +from conftest import DEFAULT_URI # type: ignore from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( @@ -13,7 +15,7 @@ from langgraph.checkpoint.postgres import PostgresSaver class TestPostgresSaver: @pytest.fixture(autouse=True) - def setup(self): + def setup(self) -> None: # objects for test setup self.config_1: RunnableConfig = { "configurable": { @@ -58,7 +60,7 @@ class TestPostgresSaver: with PostgresSaver.from_conn_string(DEFAULT_URI) as saver: saver.setup() - def test_search(self): + def test_search(self) -> None: with PostgresSaver.from_conn_string(DEFAULT_URI) as saver: # save checkpoints saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {}) @@ -66,13 +68,13 @@ class TestPostgresSaver: saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {}) # call method / assertions - query_1: CheckpointMetadata = {"source": "input"} # search by 1 key - query_2: CheckpointMetadata = { + query_1 = {"source": "input"} # search by 1 key + query_2 = { "step": 1, "writes": {"foo": "bar"}, } # search by multiple keys - query_3: CheckpointMetadata = {} # search by no keys, return all checkpoints - query_4: CheckpointMetadata = {"source": "update", "step": 1} # no match + query_3: dict[str, Any] = {} # search by no keys, return all checkpoints + query_4 = {"source": "update", "step": 1} # no match search_results_1 = list(saver.list(None, filter=query_1)) assert len(search_results_1) == 1