adding typehints

bin/parallel_sync

@@ -409,7 +409,7 @@ def run_task(
     ),
     default="multiprocess_async",
 )
-def main(config: str, nprocs: int, mode: str, verbose: bool):
+def main(config: str, nprocs: int, mode: str, verbose: bool) -> None:
     """
     TODO:
     - Track progress across cpus/threads

examples/airbnb/data.py

@@ -17,7 +17,7 @@
     help="Schema config",
     type=click.Path(exists=True),
 )
-def main(config):
+def main(config: str) -> None:
     config: str = get_config(config)
     teardown(drop_db=False, config=config)
     doc: dict = next(config_loader(config))

examples/airbnb/schema.py

@@ -118,7 +118,7 @@ def setup(config: str) -> None:
     help="Schema config",
     type=click.Path(exists=True),
 )
-def main(config):
+def main(config: str) -> None:
     config: str = get_config(config)
     teardown(config=config)
     setup(config)

examples/ancestry/data.py

@@ -14,7 +14,7 @@
     help="Schema config",
     type=click.Path(exists=True),
 )
-def main(config):
+def main(config: str) -> None:
     config: str = get_config(config)
     teardown(drop_db=False, config=config)
     doc: dict = next(config_loader(config))

examples/ancestry/schema.py

@@ -82,7 +82,7 @@ def setup(config: str) -> None:
     help="Schema config",
     type=click.Path(exists=True),
 )
-def main(config):
+def main(config: str) -> None:
     config: str = get_config(config)
     teardown(config=config)
     setup(config)

examples/book/benchmark.py

@@ -134,7 +134,7 @@ def truncate_op(session: sessionmaker, model, nsize: int) -> None:
         case_sensitive=False,
     ),
 )
-def main(config, nsize, daemon, tg_op):
+def main(config: str, nsize: int, daemon: bool, tg_op: str):
     show_settings()
 
     config: str = get_config(config)
@@ -145,7 +145,7 @@ def main(config, nsize, daemon, tg_op):
         session = Session()
 
         model = Book
-        func = {
+        func: dict = {
             INSERT: insert_op,
             UPDATE: update_op,
             DELETE: delete_op,

examples/book/data.py

@@ -36,7 +36,7 @@
     type=click.Path(exists=True),
 )
 @click.option("--nsize", "-n", default=1, help="Number of dummy data samples")
-def main(config, nsize):
+def main(config: str, nsize: int):
     config: str = get_config(config)
     teardown(drop_db=False, config=config)
 

examples/book/schema.py

@@ -259,7 +259,7 @@ def setup(config: str) -> None:
     help="Schema config",
     type=click.Path(exists=True),
 )
-def main(config):
+def main(config: str) -> None:
     config: str = get_config(config)
     teardown(config=config)
     setup(config)

examples/node/data.py

@@ -16,10 +16,10 @@
     help="Schema config",
     type=click.Path(exists=True),
 )
-def main(config):
+def main(config: str) -> None:
     config: str = get_config(config)
     teardown(drop_db=False, config=config)
-    doc = next(config_loader(config))
+    doc: dict = next(config_loader(config))
     database: str = doc.get("database", doc["index"])
     with pg_engine(database) as engine:
         Session = sessionmaker(bind=engine, autoflush=True)

examples/node/schema.py

@@ -39,7 +39,7 @@ def setup(config: str) -> None:
     help="Schema config",
     type=click.Path(exists=True),
 )
-def main(config):
+def main(config: str) -> None:
     config = get_config(config)
     teardown(config=config)
     setup(config)

examples/quiz/data.py

@@ -1,3 +1,5 @@
+import typing as t
+
 import click
 from schema import Answer, Category, PossibleAnswer, Question, RealAnswer
 from sqlalchemy.orm import sessionmaker
@@ -14,17 +16,17 @@
     help="Schema config",
     type=click.Path(exists=True),
 )
-def main(config):
+def main(config: str) -> None:
     config: str = get_config(config)
     teardown(drop_db=False, config=config)
-    doc = next(config_loader(config))
+    doc: dict = next(config_loader(config))
     database: str = doc.get("database", doc["index"])
     with pg_engine(database) as engine:
         Session = sessionmaker(bind=engine, autoflush=True)
         session = Session()
 
         # Bootstrap
-        categories = [
+        categories: t.List[Category] = [
             Category(
                 id=1,
                 uid="c001",
@@ -39,7 +41,7 @@ def main(config):
         with subtransactions(session):
             session.add_all(categories)
 
-        questions = [
+        questions: t.List[Question] = [
             Question(
                 id=1,
                 uid="q001",
@@ -58,7 +60,7 @@ def main(config):
         with subtransactions(session):
             session.add_all(questions)
 
-        answers = [
+        answers: t.List[Answer] = [
             Answer(id=1, uid="a001", text="Red"),
             Answer(id=2, uid="a002", text="Yes"),
             Answer(id=3, uid="a003", text="Green"),
@@ -67,7 +69,7 @@ def main(config):
         with subtransactions(session):
             session.add_all(answers)
 
-        possible_answers = [
+        possible_answers: t.List[PossibleAnswer] = [
             PossibleAnswer(
                 question_id=1,
                 question_uid="q001",
@@ -96,7 +98,7 @@ def main(config):
         with subtransactions(session):
             session.add_all(possible_answers)
 
-        real_answers = [
+        real_answers: t.List[RealAnswer] = [
             RealAnswer(
                 question_id=1,
                 question_uid="q001",

examples/quiz/schema.py

@@ -122,7 +122,7 @@ def setup(config: str) -> None:
     help="Schema config",
     type=click.Path(exists=True),
 )
-def main(config):
+def main(config: str) -> None:
     config: str = get_config(config)
     teardown(config=config)
     setup(config)

examples/schemas/data.py

@@ -14,10 +14,10 @@
     help="Schema config",
     type=click.Path(exists=True),
 )
-def main(config):
+def main(config: str) -> None:
     config: str = get_config(config)
     teardown(drop_db=False, config=config)
-    doc = next(config_loader(config))
+    doc: dict = next(config_loader(config))
     database: str = doc.get("database", doc["index"])
     with pg_engine(database) as engine:
         Session = sessionmaker(bind=engine, autoflush=True)

examples/schemas/schema.py

@@ -57,7 +57,7 @@ def setup(config: str) -> None:
     help="Schema config",
     type=click.Path(exists=True),
 )
-def main(config):
+def main(config: str) -> None:
     config: str = get_config(config)
     teardown(config=config)
     setup(config)

examples/social/data.py

@@ -1,3 +1,5 @@
+import typing as t
+
 import click
 from schema import Comment, Post, PostComment, Tag, User, UserPost, UserTag
 from sqlalchemy.orm import sessionmaker
@@ -14,7 +16,7 @@
     help="Schema config",
     type=click.Path(exists=True),
 )
-def main(config):
+def main(config: str) -> None:
     config: str = get_config(config)
     teardown(drop_db=False, config=config)
     doc: dict = next(config_loader(config))
@@ -24,23 +26,23 @@ def main(config):
         session = Session()
 
         # Bootstrap
-        users = [
+        users: t.List[User] = [
             User(name="Carla Ferreira Cardoso", age=19, gender="female"),
             User(name="Uwe Fuerst", age=58, gender="male"),
             User(name="Otitodilinna Chigolum", age=36, gender="male"),
         ]
         with subtransactions(session):
             session.add_all(users)
 
-        posts = [
+        posts: t.List[Post] = [
             Post(slug="post_1", title="This is the first post"),
             Post(slug="post_2", title="This is the second post"),
             Post(slug="post_3", title="This is the third post"),
         ]
         with subtransactions(session):
             session.add_all(posts)
 
-        comments = [
+        comments: t.List[Comment] = [
             Comment(
                 title="Comment 1",
                 content="This is a sample comment for comment 1",
@@ -69,7 +71,7 @@ def main(config):
         with subtransactions(session):
             session.add_all(comments)
 
-        tags = [
+        tags: t.List[Tag] = [
             Tag(name="Economics"),
             Tag(name="Career"),
             Tag(name="Political"),
@@ -86,7 +88,7 @@ def main(config):
         with subtransactions(session):
             session.add_all(tags)
 
-        user_posts = [
+        user_posts: t.List[UserPost] = [
             UserPost(
                 user=users[0],
                 post=posts[0],
@@ -103,7 +105,7 @@ def main(config):
         with subtransactions(session):
             session.add_all(user_posts)
 
-        user_tags = [
+        user_tags: t.List[UserTag] = [
             UserTag(
                 user=users[0],
                 tag=tags[0],
@@ -140,7 +142,7 @@ def main(config):
         with subtransactions(session):
             session.add_all(user_tags)
 
-        post_comments = [
+        post_comments: t.List[PostComment] = [
             PostComment(
                 post=posts[0],
                 comment=comments[0],

examples/social/schema.py

@@ -109,7 +109,7 @@ def setup(config: str) -> None:
     help="Schema config",
     type=click.Path(exists=True),
 )
-def main(config):
+def main(config: str) -> None:
     config: str = get_config(config)
     teardown(config=config)
     setup(config)

examples/starcraft/data.py

@@ -1,3 +1,5 @@
+import typing as t
+
 import click
 from schema import Specie, Structure, Unit
 from sqlalchemy.orm import sessionmaker
@@ -14,7 +16,7 @@
     help="Schema config",
     type=click.Path(exists=True),
 )
-def main(config):
+def main(config: str) -> None:
     config: str = get_config(config)
     teardown(drop_db=False, config=config)
     doc = next(config_loader(config))
@@ -23,15 +25,15 @@ def main(config):
         Session = sessionmaker(bind=engine, autoflush=True)
         session = Session()
 
-        species = [
+        species: t.List[Specie] = [
             Specie(id=1, name="Protos"),
             Specie(id=2, name="Zerg"),
             Specie(id=3, name="Terran"),
         ]
         with subtransactions(session):
             session.add_all(species)
 
-        units = [
+        units: t.List[Unit] = [
             Unit(
                 name="Archon",
                 details="Created by merging two templar units, the archon is a powerful melee unit with a very durable force shield and a strong energy-based attack.",
@@ -81,7 +83,7 @@ def main(config):
         with subtransactions(session):
             session.add_all(units)
 
-        structures = [
+        structures: t.List[Structure] = [
             Structure(
                 name="Assimilator",
                 details="Allows probes to harvest vespene gas from geysers.",

examples/starcraft/schema.py

@@ -64,7 +64,7 @@ def setup(config: str) -> None:
     help="Schema config",
     type=click.Path(exists=True),
 )
-def main(config):
+def main(config: str) -> None:
     config: str = get_config(config)
     teardown(config=config)
     setup(config)

examples/through/data.py

@@ -1,3 +1,5 @@
+import typing as t
+
 import click
 from schema import Customer, CustomerGroup, Group
 from sqlalchemy.orm import sessionmaker
@@ -15,7 +17,7 @@
     help="Schema config",
     type=click.Path(exists=True),
 )
-def main(config):
+def main(config: str) -> None:
     config: str = get_config(config)
     teardown(drop_db=False, config=config)
 
@@ -29,23 +31,23 @@ def main(config):
             Session = sessionmaker(bind=connection, autoflush=True)
             session = Session()
 
-            customers = [
+            customers: t.List[Customer] = [
                 Customer(name="CustomerA"),
                 Customer(name="CustomerB"),
                 Customer(name="CustomerC"),
             ]
             with subtransactions(session):
                 session.add_all(customers)
 
-            groups = [
+            groups: t.List[Group] = [
                 Group(group_name="GroupA"),
                 Group(group_name="GroupB"),
                 Group(group_name="GroupC"),
             ]
             with subtransactions(session):
                 session.add_all(groups)
 
-            customers_groups = [
+            customers_groups: t.List[CustomerGroup] = [
                 CustomerGroup(customer=customers[0], group=groups[0]),
                 CustomerGroup(customer=customers[1], group=groups[1]),
                 CustomerGroup(customer=customers[2], group=groups[2]),

examples/through/schema.py

@@ -83,7 +83,7 @@ def setup(config: str) -> None:
     help="Schema config",
     type=click.Path(exists=True),
 )
-def main(config):
+def main(config: str) -> None:
     config: str = get_config(config)
     teardown(config=config)
     setup(config)

plugins/openai_plugin.py

@@ -17,7 +17,7 @@ def __init__(self) -> None:
         self.client: OpenAI = OpenAI()
         self.model: str = "text-embedding-3-small"
         # vector dims must match models input dims
-        self.vector_dims = 1536
+        self.vector_dims: int = 1536
 
     name: str = "TextEmbedding3Small"