add better python annotations (#192)

crates/python/python/cryo/__init__.py

@@ -2,7 +2,7 @@
 
 # ruff: noqa: F401
 
-from ._cryo_rust import __version__
+from ._cryo_rust import __version__  # type: ignore
 from ._freeze import async_freeze
 from ._freeze import freeze
 from ._collect import async_collect

crates/python/python/cryo/_args.py

@@ -12,7 +12,6 @@ def parse_cli_args(
     start_block: _spec.BlockReference | None = None,
     end_block: _spec.BlockReference | None = None,
     file_format: _spec.FileFormat = 'parquet',
-    verbose: bool = True,
     **kwargs: Unpack[_spec.CryoCliArgs],
 ) -> _spec.CryoCliArgs:
     """
@@ -37,7 +36,7 @@ def parse_cli_args(
     else:
         raise Exception('unknown file_format')
 
-    kwargs['no_verbose'] = not verbose
+    kwargs['no_verbose'] = not kwargs.get('verbose', True)
 
     return kwargs
 

crates/python/python/cryo/_collect.py

@@ -3,14 +3,51 @@
 
 if typing.TYPE_CHECKING:
     from typing_extensions import Unpack
-    from typing import Any
+    from typing import Any, Literal, TypeVar
 
     import pandas as pd
     import polars as pl
     from . import _spec
 
     ListOfDicts = list[dict[str, Any]]
     DictOfLists = dict[str, list[Any]]
+    T = TypeVar('T', pl.DataFrame, pd.DataFrame, ListOfDicts, DictOfLists)
+
+
+@typing.overload
+async def async_collect(
+    datatype: _spec.Datatype,
+    output_format: Literal['polars'] = 'polars',
+    **kwargs: Unpack[_spec.CryoCliArgs],
+) -> pl.DataFrame:
+    ...
+
+
+@typing.overload
+async def async_collect(
+    datatype: _spec.Datatype,
+    output_format: Literal['pandas'],
+    **kwargs: Unpack[_spec.CryoCliArgs],
+) -> pd.DataFrame:
+    ...
+
+
+@typing.overload
+async def async_collect(
+    datatype: _spec.Datatype,
+    output_format: Literal['list'],
+    **kwargs: Unpack[_spec.CryoCliArgs],
+) -> ListOfDicts:
+    ...
+
+
+@typing.overload
+async def async_collect(
+    datatype: _spec.Datatype,
+    output_format: Literal['dict'],
+    **kwargs: Unpack[_spec.CryoCliArgs],
+) -> DictOfLists:
+    ...
 
 
 async def async_collect(
@@ -45,6 +82,42 @@ async def async_collect(
         raise Exception('unknown output format')
 
 
+@typing.overload
+def collect(
+    datatype: _spec.Datatype,
+    output_format: Literal['polars'] = 'polars',
+    **kwargs: Unpack[_spec.CryoCliArgs],
+) -> pl.DataFrame:
+    ...
+
+
+@typing.overload
+def collect(
+    datatype: _spec.Datatype,
+    output_format: Literal['pandas'],
+    **kwargs: Unpack[_spec.CryoCliArgs],
+) -> pd.DataFrame:
+    ...
+
+
+@typing.overload
+def collect(
+    datatype: _spec.Datatype,
+    output_format: Literal['list'],
+    **kwargs: Unpack[_spec.CryoCliArgs],
+) -> ListOfDicts:
+    ...
+
+
+@typing.overload
+def collect(
+    datatype: _spec.Datatype,
+    output_format: Literal['dict'],
+    **kwargs: Unpack[_spec.CryoCliArgs],
+) -> DictOfLists:
+    ...
+
+
 def collect(
     datatype: _spec.Datatype,
     output_format: _spec.PythonOutput = 'polars',
@@ -58,11 +131,14 @@ def collect(
 
     try:
         import concurrent.futures
+
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
         with concurrent.futures.ThreadPoolExecutor() as executor:
             future = executor.submit(loop.run_until_complete, coroutine)  # type: ignore
-            return future.result()  # type: ignore
+            result: T = future.result()  # type: ignore
     except RuntimeError:
-        return asyncio.run(coroutine)
+        result = asyncio.run(coroutine)
+
+    return result
 

crates/python/python/cryo/_spec.py

@@ -47,11 +47,33 @@ class CryoCliArgs(TypedDict, total=False):
         n_row_groups: int | None
         no_stats: bool
         compression: str | None
-        contract: str | bytes | None
-        topic0: str | bytes | None
-        topic1: str | bytes | None
-        topic2: str | bytes | None
-        topic3: str | bytes | None
+        contract: typing.Sequence[str | bytes | None]
+        topic0: typing.Sequence[str | bytes | None]
+        topic1: typing.Sequence[str | bytes | None]
+        topic2: typing.Sequence[str | bytes | None]
+        topic3: typing.Sequence[str | bytes | None]
         inner_request_size: int | None
         no_verbose: bool
 
+        timestamps: typing.Sequence[str] | None
+        txs: typing.Sequence[str] | None
+        u256_types: typing.Sequence[str] | None
+        exclude_failed: bool
+        chunk_order: str | None
+        max_retries: int
+        initial_backoff: int
+        partition_by: typing.Sequence[str] | None
+        subdirs: typing.Sequence[str]
+        label: str | None
+        report_dir: str | None
+        no_report: bool
+        address: typing.Sequence[str] | None
+        to_address: typing.Sequence[str] | None
+        from_address: typing.Sequence[str] | None
+        call_data: typing.Sequence[str] | None
+        function: typing.Sequence[str] | None
+        inputs: typing.Sequence[str] | None
+        slot: typing.Sequence[str] | None
+        js_tracer: str | None
+        verbose: bool
+        event_signature: str | None