refactor: add type annotations and clean up code

ddss/dump.py

@@ -1,10 +1,11 @@
 import asyncio
 from sqlalchemy import select
+from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession
 from apyds_bnf import unparse
 from .orm import Facts, Ideas
 
 
-async def main(session):
+async def main(session: async_sessionmaker[AsyncSession]) -> None:
     try:
         async with session() as sess:
             for i in await sess.scalars(select(Ideas)):

ddss/dump.ts

@@ -1,12 +1,8 @@
 import type { Sequelize } from "sequelize";
 import { unparse } from "atsds-bnf";
-import { Fact, Idea, initializeDatabase } from "./orm.ts";
-
-export async function main(addr: string, sequelize?: Sequelize) {
-    if (!sequelize) {
-        sequelize = await initializeDatabase(addr);
-    }
+import { Fact, Idea } from "./orm.ts";
 
+export async function main(sequelize: Sequelize): Promise<void> {
     const ideas = await Idea.findAll();
     for (const idea of ideas) {
         console.log("idea:", unparse(idea.data));

ddss/egg.py

@@ -1,11 +1,12 @@
 import asyncio
 from sqlalchemy import select
+from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession
 from apyds import Rule
 from .orm import insert_or_ignore, Facts, Ideas
 from .egraph import Search
 
 
-async def main(session):
+async def main(session: async_sessionmaker[AsyncSession]) -> None:
     try:
         search = Search()
         pool = []

ddss/egg.ts

@@ -1,9 +1,9 @@
 import { Op, type Sequelize } from "sequelize";
 import { Rule } from "atsds";
 import { Search } from "./egraph.ts";
-import { Fact, Idea, initializeDatabase, insertOrIgnore } from "./orm.ts";
+import { Fact, Idea, insertOrIgnore } from "./orm.ts";
 
-export async function main(sequelize: Sequelize) {
+export async function main(sequelize: Sequelize): Promise<void> {
     const search = new Search();
     let pool: Rule[] = [];
     let maxFact = -1;

ddss/egraph.ts

@@ -9,15 +9,16 @@ function extractLhsRhsFromRule(data: Rule): [Term, Term] | null {
     if (data.length() !== 0) {
         return null;
     }
-    const term = data.conclusion();
-    const inner = term.term();
-    if (!(inner instanceof List)) {
+    const term = data.conclusion().term();
+    if (!(term instanceof List)) {
         return null;
     }
-    if (!(inner.length() === 4 && inner.getitem(0).toString() === "binary" && inner.getitem(1).toString() === "==")) {
+    if (!(term.length() === 4 && term.getitem(0).toString() === "binary" && term.getitem(1).toString() === "==")) {
         return null;
     }
-    return [inner.getitem(2), inner.getitem(3)];
+    const lhs = term.getitem(2);
+    const rhs = term.getitem(3);
+    return [lhs, rhs];
 }
 
 function buildLhsRhsToTerm(lhs: Term, rhs: Term): Term {

ddss/input.py

@@ -1,12 +1,13 @@
 import asyncio
 from prompt_toolkit import PromptSession
 from prompt_toolkit.patch_stdout import patch_stdout
+from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession
 from apyds_bnf import parse
 from .orm import insert_or_ignore, Facts, Ideas
 from .utility import str_rule_get_str_idea
 
 
-async def main(session):
+async def main(session: async_sessionmaker[AsyncSession]) -> None:
     try:
         prompt = PromptSession()
         while True:
@@ -17,9 +18,7 @@ async def main(session):
                 raise asyncio.CancelledError()
 
             data = line.strip()
-            if data == "":
-                continue
-            if data.startswith("//"):
+            if data == "" or data.startswith("//"):
                 continue
 
             try:

ddss/input.ts

@@ -2,10 +2,10 @@ import * as readline from "node:readline/promises";
 import { stdin as input, stdout as output } from "node:process";
 import type { Sequelize } from "sequelize";
 import { parse } from "atsds-bnf";
-import { Fact, Idea, initializeDatabase, insertOrIgnore } from "./orm.ts";
+import { Fact, Idea, insertOrIgnore } from "./orm.ts";
 import { patchStdout, strRuleGetStrIdea } from "./utility.ts";
 
-export async function main(sequelize: Sequelize) {
+export async function main(sequelize: Sequelize): Promise<void> {
     const rl = readline.createInterface({ input, output });
     rl.setPrompt("input: ");
     const unpatch = patchStdout(rl);
@@ -21,10 +21,9 @@ export async function main(sequelize: Sequelize) {
 
         try {
             const ds = parse(data);
-            const dsStr = ds.toString();
 
-            await insertOrIgnore(Fact, dsStr);
-            const idea = strRuleGetStrIdea(dsStr);
+            await insertOrIgnore(Fact, ds);
+            const idea = strRuleGetStrIdea(ds);
             if (idea) {
                 await insertOrIgnore(Idea, idea);
             }

ddss/load.py

@@ -1,18 +1,17 @@
 import asyncio
 import sys
+from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession
 from apyds_bnf import parse
 from .orm import insert_or_ignore, Facts, Ideas
 from .utility import str_rule_get_str_idea
 
 
-async def main(session):
+async def main(session: async_sessionmaker[AsyncSession]) -> None:
     try:
         async with session() as sess:
             for line in sys.stdin:
                 data = line.strip()
-                if data == "":
-                    continue
-                if data.startswith("//"):
+                if data == "" or data.startswith("//"):
                     continue
 
                 try:

ddss/load.ts

@@ -2,10 +2,10 @@ import * as readline from "node:readline/promises";
 import { stdin as input } from "node:process";
 import type { Sequelize } from "sequelize";
 import { parse } from "atsds-bnf";
-import { Fact, Idea, initializeDatabase, insertOrIgnore } from "./orm.ts";
+import { Fact, Idea, insertOrIgnore } from "./orm.ts";
 import { strRuleGetStrIdea } from "./utility.ts";
 
-export async function main(sequelize: Sequelize) {
+export async function main(sequelize: Sequelize): Promise<void> {
     const rl = readline.createInterface({
         input,
         terminal: false,
@@ -19,10 +19,9 @@ export async function main(sequelize: Sequelize) {
 
         try {
             const ds = parse(data);
-            const dsStr = ds.toString();
 
-            await insertOrIgnore(Fact, dsStr);
-            const idea = strRuleGetStrIdea(dsStr);
+            await insertOrIgnore(Fact, ds);
+            const idea = strRuleGetStrIdea(ds);
             if (idea) {
                 await insertOrIgnore(Idea, idea);
             }

ddss/main.py

@@ -1,8 +1,9 @@
 import asyncio
 import tempfile
 import pathlib
-from typing import Annotated, Optional
+from typing import Annotated, Optional, Awaitable
 import tyro
+from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession
 from .orm import initialize_database
 from .search import main as search
 from .egg import main as egg
@@ -11,7 +12,7 @@
 from .load import main as load
 from .dump import main as dump
 
-component_map = {
+component_map: dict[str, callable[[async_sessionmaker[AsyncSession]], Awaitable[None]]] = {
     "search": search,
     "egg": egg,
     "input": input,
@@ -29,7 +30,7 @@ async def run(addr: str, components: list[str]) -> None:
             coroutines = [component_map[component](session) for component in components]
         except KeyError as e:
             print(f"error: unsupported component: {str(e)}")
-            raise asyncio.CancelledError()
+            return
 
         await asyncio.wait(
             [asyncio.create_task(coro) for coro in coroutines],

ddss/main.ts

@@ -11,7 +11,7 @@ import { main as load } from "./load.ts";
 import { main as output } from "./output.ts";
 import { initializeDatabase } from "./orm.ts";
 
-type ComponentMain = (addr: string, sequelize: Sequelize) => Promise<void>;
+type ComponentMain = (sequelize: Sequelize) => Promise<void>;
 
 const componentMap: Record<string, ComponentMain> = {
     search,
@@ -22,27 +22,27 @@ const componentMap: Record<string, ComponentMain> = {
     dump,
 };
 
-async function run(addr: string, components: string[]) {
+async function run(addr: string, components: string[]): Promise<void> {
+    const sequelize = await initializeDatabase(addr);
+
     for (const name of components) {
         if (!(name in componentMap)) {
             console.error(`error: unsupported component: ${name}`);
             return;
         }
     }
 
-    const sequelize = await initializeDatabase(addr);
-
     const promises = components.map((name) => {
         const component = componentMap[name]!;
-        return component(addr, sequelize);
+        return component(sequelize);
     });
 
     await Promise.race(promises);
 
     await sequelize.close();
 }
 
-export function cli() {
+export async function cli(): Promise<void> {
     const program = new Command();
 
     program

ddss/orm.ts

@@ -5,7 +5,6 @@ import {
     type CreationOptional,
     type InferAttributes,
     type InferCreationAttributes,
-    type ModelStatic,
 } from "sequelize";
 
 class Fact extends Model<InferAttributes<Fact>, InferCreationAttributes<Fact>> {

ddss/output.py

@@ -1,10 +1,11 @@
 import asyncio
 from sqlalchemy import select
+from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession
 from apyds_bnf import unparse
 from .orm import Facts, Ideas
 
 
-async def main(session):
+async def main(session: async_sessionmaker[AsyncSession]) -> None:
     try:
         max_fact = -1
         max_idea = -1

ddss/output.ts

@@ -1,8 +1,8 @@
 import { Op, type Sequelize } from "sequelize";
 import { unparse } from "atsds-bnf";
-import { Fact, Idea, initializeDatabase } from "./orm.ts";
+import { Fact, Idea } from "./orm.ts";
 
-export async function main(sequelize: Sequelize) {
+export async function main(sequelize: Sequelize): Promise<void> {
     let maxFact = -1;
     let maxIdea = -1;
 

ddss/search.py

@@ -1,11 +1,12 @@
 import asyncio
 from sqlalchemy import select
+from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession
 from apyds import Search
 from .orm import insert_or_ignore, Facts, Ideas
 from .utility import str_rule_get_str_idea
 
 
-async def main(session):
+async def main(session: async_sessionmaker[AsyncSession]) -> None:
     try:
         search = Search()
         max_fact = -1

ddss/search.ts

@@ -1,10 +1,10 @@
 import { Op, type Sequelize } from "sequelize";
 import { Search } from "atsds";
 import type { Rule } from "atsds";
-import { Fact, Idea, initializeDatabase, insertOrIgnore } from "./orm.ts";
+import { Fact, Idea, insertOrIgnore } from "./orm.ts";
 import { strRuleGetStrIdea } from "./utility.ts";
 
-export async function main(sequelize: Sequelize) {
+export async function main(sequelize: Sequelize): Promise<void> {
     const search = new Search();
     let maxFact = -1;
 
@@ -21,9 +21,9 @@ export async function main(sequelize: Sequelize) {
         const tasks: Promise<void>[] = [];
 
         const handler = (rule: Rule) => {
-            const dsStr = rule.toString();
-            tasks.push(insertOrIgnore(Fact, dsStr));
-            const idea = strRuleGetStrIdea(dsStr);
+            const ds = rule.toString();
+            tasks.push(insertOrIgnore(Fact, ds));
+            const idea = strRuleGetStrIdea(ds);
             if (idea) {
                 tasks.push(insertOrIgnore(Idea, idea));
             }

ddss/utility.ts

@@ -10,7 +10,7 @@ export function strRuleGetStrIdea(data: string): string | null {
     return null;
 }
 
-export function patchStdout(rl: Interface) {
+export function patchStdout(rl: Interface): () => void {
     const originalWrite = stdout.write;
     const originalLog = console.log;
     const originalError = console.error;