Merge pull request #82 from simonsobs/dev

Fix all mypy errors with the `disallow_untyped_defs` option

pyproject.toml

@@ -93,6 +93,7 @@ exclude = '''(?x)(
     src/nextline_rdb/alembic/versions/.*\.py$
     | example_script/.*\.py$
 )'''
+disallow_untyped_defs = true
 
 [[tool.mypy.overrides]]
 module = ["dynaconf.*", "async_asgi_testclient.*", "apluggy.*"]

src/nextline_rdb/alembic/models/rev_4dc6a93dfed8/model_script.py

@@ -29,7 +29,7 @@ def current(self) -> bool:
         return self._current is not None
 
     @current.setter
-    def current(self, value: bool):
+    def current(self, value: bool) -> None:
         if value and self._current is None:
             self._current = CurrentScript(script=self)
         elif not value:

src/nextline_rdb/alembic/models/rev_4dc6a93dfed8/tests/test_model_prompt.py

@@ -12,7 +12,7 @@
 
 
 @given(st.data())
-async def test_repr(data: st.DataObject):
+async def test_repr(data: st.DataObject) -> None:
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
             model = data.draw(st_model_prompt())

src/nextline_rdb/alembic/models/rev_4dc6a93dfed8/tests/test_model_run.py

@@ -10,7 +10,7 @@
 
 
 @given(run=st_model_run(generate_traces=False))
-async def test_repr(run: Run):
+async def test_repr(run: Run) -> None:
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
             session.add(run)
@@ -24,7 +24,7 @@ async def test_repr(run: Run):
 
 
 @given(run=st_model_run(generate_traces=True))
-async def test_cascade(run: Run):
+async def test_cascade(run: Run) -> None:
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
             session.add(run)

src/nextline_rdb/alembic/models/rev_4dc6a93dfed8/tests/test_model_stdout.py

@@ -12,7 +12,7 @@
 
 
 @given(st.data())
-async def test_repr(data: st.DataObject):
+async def test_repr(data: st.DataObject) -> None:
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
             model = data.draw(st_model_stdout())

src/nextline_rdb/alembic/models/rev_4dc6a93dfed8/tests/test_model_trace.py

@@ -12,7 +12,7 @@
 
 
 @given(st.data())
-async def test_repr(data: st.DataObject):
+async def test_repr(data: st.DataObject) -> None:
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
             model = data.draw(st_model_trace())

src/nextline_rdb/alembic/models/rev_4dc6a93dfed8/tests/test_repr_val.py

@@ -45,7 +45,7 @@ def test_enum(value: Color) -> None:
 
 
 @st.composite
-def st_enum_type(draw: st.DrawFn):
+def st_enum_type(draw: st.DrawFn) -> type[Enum]:
     '''Generate an Enum type.
 
     >>> enum_type = st_enum_type().example()
@@ -59,11 +59,11 @@ def st_enum_type(draw: st.DrawFn):
     names_ = st.text(ascii_lowercase, min_size=1)
     names = st.builds(lambda x: x.capitalize(), names_).filter(str.isidentifier)
     values = st.lists(st.text(ascii_uppercase, min_size=1), min_size=1, unique=True)
-    return draw(st.builds(Enum, names, values))
+    return draw(st.builds(Enum, names, values))  # type: ignore
 
 
 @given(st.data())
-def test_arbitrary_enum(data: st.DataObject):
+def test_arbitrary_enum(data: st.DataObject) -> None:
     enum_type = data.draw(st_enum_type())
     item = data.draw(st.sampled_from(enum_type))
     repr_ = repr_val(item)

src/nextline_rdb/alembic/models/rev_5a08750d6760/tests/test_model_prompt.py

@@ -12,7 +12,7 @@
 
 
 @given(st.data())
-async def test_repr(data: st.DataObject):
+async def test_repr(data: st.DataObject) -> None:
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
             model = data.draw(st_model_prompt())

src/nextline_rdb/alembic/models/rev_5a08750d6760/tests/test_model_run.py

@@ -11,7 +11,7 @@
 
 
 @given(st.data())
-async def test_repr(data: st.DataObject):
+async def test_repr(data: st.DataObject) -> None:
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
             model = data.draw(st_model_run(generate_traces=False))

src/nextline_rdb/alembic/models/rev_5a08750d6760/tests/test_model_stdout.py

@@ -12,7 +12,7 @@
 
 
 @given(st.data())
-async def test_repr(data: st.DataObject):
+async def test_repr(data: st.DataObject) -> None:
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
             model = data.draw(st_model_stdout())

src/nextline_rdb/alembic/models/rev_5a08750d6760/tests/test_model_trace.py

@@ -12,7 +12,7 @@
 
 
 @given(st.data())
-async def test_repr(data: st.DataObject):
+async def test_repr(data: st.DataObject) -> None:
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
             model = data.draw(st_model_trace())

src/nextline_rdb/alembic/models/rev_5a08750d6760/tests/test_repr_val.py

@@ -45,7 +45,7 @@ def test_enum(value: Color) -> None:
 
 
 @st.composite
-def st_enum_type(draw: st.DrawFn):
+def st_enum_type(draw: st.DrawFn) -> type[Enum]:
     '''Generate an Enum type.
 
     >>> enum_type = st_enum_type().example()
@@ -59,11 +59,11 @@ def st_enum_type(draw: st.DrawFn):
     names_ = st.text(ascii_lowercase, min_size=1)
     names = st.builds(lambda x: x.capitalize(), names_).filter(str.isidentifier)
     values = st.lists(st.text(ascii_uppercase, min_size=1), min_size=1, unique=True)
-    return draw(st.builds(Enum, names, values))
+    return draw(st.builds(Enum, names, values))  # type: ignore
 
 
 @given(st.data())
-def test_arbitrary_enum(data: st.DataObject):
+def test_arbitrary_enum(data: st.DataObject) -> None:
     enum_type = data.draw(st_enum_type())
     item = data.draw(st.sampled_from(enum_type))
     repr_ = repr_val(item)

src/nextline_rdb/alembic/models/rev_f3edea6dbde2/model_script.py

@@ -29,7 +29,7 @@ def current(self) -> bool:
         return self._current is not None
 
     @current.setter
-    def current(self, value: bool):
+    def current(self, value: bool) -> None:
         if value and self._current is None:
             self._current = CurrentScript(script=self)
         elif not value:

src/nextline_rdb/alembic/models/rev_f3edea6dbde2/tests/test_model_prompt.py

@@ -12,7 +12,7 @@
 
 
 @given(st.data())
-async def test_repr(data: st.DataObject):
+async def test_repr(data: st.DataObject) -> None:
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
             model = data.draw(st_model_prompt())

src/nextline_rdb/alembic/models/rev_f3edea6dbde2/tests/test_model_run.py

@@ -10,7 +10,7 @@
 
 
 @given(run=st_model_run(generate_traces=False))
-async def test_repr(run: Run):
+async def test_repr(run: Run) -> None:
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
             session.add(run)
@@ -25,7 +25,7 @@ async def test_repr(run: Run):
 
 @settings(phases=(Phase.generate,))  # Avoid shrinking
 @given(run=st_model_run(generate_traces=True))
-async def test_cascade(run: Run):
+async def test_cascade(run: Run) -> None:
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
             session.add(run)

src/nextline_rdb/alembic/models/rev_f3edea6dbde2/tests/test_model_stdout.py

@@ -12,7 +12,7 @@
 
 
 @given(st.data())
-async def test_repr(data: st.DataObject):
+async def test_repr(data: st.DataObject) -> None:
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
             model = data.draw(st_model_stdout())

src/nextline_rdb/alembic/models/rev_f3edea6dbde2/tests/test_model_trace.py

@@ -12,7 +12,7 @@
 
 
 @given(st.data())
-async def test_repr(data: st.DataObject):
+async def test_repr(data: st.DataObject) -> None:
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
             model = data.draw(st_model_trace())

src/nextline_rdb/alembic/models/rev_f3edea6dbde2/tests/test_repr_val.py

@@ -45,7 +45,7 @@ def test_enum(value: Color) -> None:
 
 
 @st.composite
-def st_enum_type(draw: st.DrawFn):
+def st_enum_type(draw: st.DrawFn) -> type[Enum]:
     '''Generate an Enum type.
 
     >>> enum_type = st_enum_type().example()
@@ -59,11 +59,11 @@ def st_enum_type(draw: st.DrawFn):
     names_ = st.text(ascii_lowercase, min_size=1)
     names = st.builds(lambda x: x.capitalize(), names_).filter(str.isidentifier)
     values = st.lists(st.text(ascii_uppercase, min_size=1), min_size=1, unique=True)
-    return draw(st.builds(Enum, names, values))
+    return draw(st.builds(Enum, names, values))  # type: ignore
 
 
 @given(st.data())
-def test_arbitrary_enum(data: st.DataObject):
+def test_arbitrary_enum(data: st.DataObject) -> None:
     enum_type = data.draw(st_enum_type())
     item = data.draw(st.sampled_from(enum_type))
     repr_ = repr_val(item)

src/nextline_rdb/alembic/models/rev_f9a742bb2297/tests/test_model_prompt.py

@@ -12,7 +12,7 @@
 
 
 @given(st.data())
-async def test_repr(data: st.DataObject):
+async def test_repr(data: st.DataObject) -> None:
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
             model = data.draw(st_model_prompt())

src/nextline_rdb/alembic/models/rev_f9a742bb2297/tests/test_model_run.py

@@ -10,7 +10,7 @@
 
 
 @given(run=st_model_run(generate_traces=False))
-async def test_repr(run: Run):
+async def test_repr(run: Run) -> None:
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
             session.add(run)
@@ -24,7 +24,7 @@ async def test_repr(run: Run):
 
 
 @given(run=st_model_run(generate_traces=True))
-async def test_cascade(run: Run):
+async def test_cascade(run: Run) -> None:
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
             session.add(run)

src/nextline_rdb/alembic/models/rev_f9a742bb2297/tests/test_model_stdout.py

@@ -12,7 +12,7 @@
 
 
 @given(st.data())
-async def test_repr(data: st.DataObject):
+async def test_repr(data: st.DataObject) -> None:
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
             model = data.draw(st_model_stdout())

src/nextline_rdb/alembic/models/rev_f9a742bb2297/tests/test_model_trace.py

@@ -12,7 +12,7 @@
 
 
 @given(st.data())
-async def test_repr(data: st.DataObject):
+async def test_repr(data: st.DataObject) -> None:
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
             model = data.draw(st_model_trace())

src/nextline_rdb/alembic/models/rev_f9a742bb2297/tests/test_repr_val.py

@@ -45,7 +45,7 @@ def test_enum(value: Color) -> None:
 
 
 @st.composite
-def st_enum_type(draw: st.DrawFn):
+def st_enum_type(draw: st.DrawFn) -> type[Enum]:
     '''Generate an Enum type.
 
     >>> enum_type = st_enum_type().example()
@@ -59,11 +59,11 @@ def st_enum_type(draw: st.DrawFn):
     names_ = st.text(ascii_lowercase, min_size=1)
     names = st.builds(lambda x: x.capitalize(), names_).filter(str.isidentifier)
     values = st.lists(st.text(ascii_uppercase, min_size=1), min_size=1, unique=True)
-    return draw(st.builds(Enum, names, values))
+    return draw(st.builds(Enum, names, values))  # type: ignore
 
 
 @given(st.data())
-def test_arbitrary_enum(data: st.DataObject):
+def test_arbitrary_enum(data: st.DataObject) -> None:
     enum_type = data.draw(st_enum_type())
     item = data.draw(st.sampled_from(enum_type))
     repr_ = repr_val(item)

src/nextline_rdb/db.py

@@ -2,7 +2,7 @@
 from logging import getLogger
 from os import PathLike
 from pathlib import Path
-from typing import Optional, Type
+from typing import Any, Optional, Type
 
 from alembic.config import Config
 from alembic.migration import MigrationContext
@@ -136,7 +136,7 @@ async def __aenter__(self) -> 'DB':
         await self.start()
         return self
 
-    async def __aexit__(self, *_, **__) -> None:
+    async def __aexit__(self, *_: Any, **__: Any) -> None:
         await self.aclose()
 
 
@@ -168,7 +168,7 @@ def do_run_migrations(connection: Connection) -> None:
 
 
 @event.listens_for(Engine, 'connect')
-def set_sqlite_pragma(dbapi_connection, connection_record):
+def set_sqlite_pragma(dbapi_connection: Any, connection_record: Any) -> None:
     '''Enable foreign key constraints in SQLite.
 
     The code copied from the SQLAlchemy documentation:

src/nextline_rdb/models/tests/test_cascade.py

@@ -13,19 +13,19 @@
 
 @settings(phases=(Phase.generate,))  # Avoid shrinking
 @given(parent=st_model_run(generate_traces=True))
-async def test_run(parent: Run):
+async def test_run(parent: Run) -> None:
     await assert_cascade(parent, Run, [Prompt, Trace, TraceCall, Stdout])
 
 
 @settings(phases=(Phase.generate,))  # Avoid shrinking
 @given(parent=st_model_trace(generate_trace_calls=True, generate_prompts=True))
-async def test_trace(parent: Trace):
+async def test_trace(parent: Trace) -> None:
     await assert_cascade(parent, Trace, [Prompt, TraceCall])
 
 
 @settings(phases=(Phase.generate,))  # Avoid shrinking
 @given(parent=st_model_trace_call(generate_prompts=True))
-async def test_trace_call(parent: TraceCall):
+async def test_trace_call(parent: TraceCall) -> None:
     await assert_cascade(parent, TraceCall, [Prompt])  # type: ignore
 
 

src/nextline_rdb/models/tests/test_repr_val.py

@@ -45,7 +45,7 @@ def test_enum(value: Color) -> None:
 
 
 @st.composite
-def st_enum_type(draw: st.DrawFn):
+def st_enum_type(draw: st.DrawFn) -> type[Enum]:
     '''Generate an Enum type.
 
     >>> enum_type = st_enum_type().example()
@@ -59,11 +59,11 @@ def st_enum_type(draw: st.DrawFn):
     names_ = st.text(ascii_lowercase, min_size=1)
     names = st.builds(lambda x: x.capitalize(), names_).filter(str.isidentifier)
     values = st.lists(st.text(ascii_uppercase, min_size=1), min_size=1, unique=True)
-    return draw(st.builds(Enum, names, values))
+    return draw(st.builds(Enum, names, values))  # type: ignore
 
 
 @given(st.data())
-def test_arbitrary_enum(data: st.DataObject):
+def test_arbitrary_enum(data: st.DataObject) -> None:
     enum_type = data.draw(st_enum_type())
     item = data.draw(st.sampled_from(enum_type))
     repr_ = repr_val(item)

src/nextline_rdb/pagination.py

@@ -1,7 +1,7 @@
 from collections.abc import Sequence
 from typing import Any, NamedTuple, Optional, Type, TypeVar
 
-from sqlalchemy import func, select
+from sqlalchemy import ScalarResult, func, select
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import DeclarativeBase, aliased, selectinload
 from sqlalchemy.sql.selectable import Select
@@ -36,7 +36,7 @@ async def load_models(
     after: Optional[_Id] = None,
     first: Optional[int] = None,
     last: Optional[int] = None,
-):
+) -> ScalarResult[T]:
     sort = sort or []
 
     if id_field not in {s.field for s in sort}:

src/nextline_rdb/schema/nodes/prompt_node.py

@@ -48,7 +48,7 @@ def trace_call(
         return TraceCallNode.from_model(self._model.trace_call)
 
     @classmethod
-    def from_model(cls: type['PromptNode'], model: db_models.Prompt):
+    def from_model(cls: type['PromptNode'], model: db_models.Prompt) -> 'PromptNode':
         return cls(
             _model=model,
             id=model.id,

src/nextline_rdb/schema/nodes/run_node.py

@@ -154,7 +154,7 @@ class RunNode:
     ] = strawberry.field(resolver=_resolve_stdouts)
 
     @classmethod
-    def from_model(cls: type['RunNode'], model: Run):
+    def from_model(cls: type['RunNode'], model: Run) -> 'RunNode':
         script = model.script.script if model.script else None
         return cls(
             _model=model,