getting close on types

This commit is contained in:
James Turk 2023-07-11 23:53:18 -05:00
parent d33483f136
commit 0998fe1e6a
3 changed files with 40 additions and 35 deletions

View File

@ -12,7 +12,7 @@ PydanticModel = Type[BaseModel]
class Beaker(abc.ABC): class Beaker(abc.ABC):
def __init__(self, name: str, model: PydanticModel, recipe: Recipe): def __init__(self, name: str, model: PydanticModel, recipe: "Recipe"):
self.name = name self.name = name
self.model = model self.model = model
self.recipe = recipe self.recipe = recipe
@ -42,7 +42,7 @@ class Beaker(abc.ABC):
class TempBeaker(Beaker): class TempBeaker(Beaker):
def __init__(self, name: str, model: PydanticModel | None, recipe: Recipe): def __init__(self, name: str, model: PydanticModel, recipe: "Recipe"):
super().__init__(name, model, recipe) super().__init__(name, model, recipe)
self._items: list[tuple[str, BaseModel]] = [] self._items: list[tuple[str, BaseModel]] = []
@ -62,11 +62,11 @@ class TempBeaker(Beaker):
class SqliteBeaker(Beaker): class SqliteBeaker(Beaker):
def __init__(self, name: str, model: PydanticModel, recipe: Recipe): def __init__(self, name: str, model: PydanticModel, recipe: "Recipe"):
super().__init__(name, model, recipe) super().__init__(name, model, recipe)
# create table if it doesn't exist # create table if it doesn't exist
self.cursor = self.recipe.db.cursor() self.cursor = self.recipe.db.cursor()
self.cursor.row_factory = sqlite3.Row self.cursor.row_factory = sqlite3.Row # type: ignore
self.cursor.execute( self.cursor.execute(
f"CREATE TABLE IF NOT EXISTS {self.name} (uuid TEXT PRIMARY KEY, data JSON)" f"CREATE TABLE IF NOT EXISTS {self.name} (uuid TEXT PRIMARY KEY, data JSON)"
) )

View File

@ -67,7 +67,7 @@ def run(
if not input and not has_data: if not input and not has_data:
typer.secho("No data; pass --input to seed beaker(s)", fg=typer.colors.RED) typer.secho("No data; pass --input to seed beaker(s)", fg=typer.colors.RED)
raise typer.Exit(1) raise typer.Exit(1)
for input_str in input: # type: ignore for input_str in input or []:
beaker, filename = input_str.split("=") beaker, filename = input_str.split("=")
ctx.obj.csv_to_beaker(filename, beaker) ctx.obj.csv_to_beaker(filename, beaker)
ctx.obj.run_once(start, end) ctx.obj.run_once(start, end)

View File

@ -5,11 +5,10 @@ import inspect
import sqlite3 import sqlite3
import hashlib import hashlib
import asyncio import asyncio
import networkx import networkx # type: ignore
from collections import defaultdict, Counter from collections import defaultdict, Counter
from dataclasses import dataclass # TODO: pydantic? from typing import Iterable, Callable, Type
from pydantic import BaseModel from pydantic import BaseModel, ConfigDict
from typing import Iterable, Callable
from structlog import get_logger from structlog import get_logger
from .beakers import Beaker, SqliteBeaker, TempBeaker from .beakers import Beaker, SqliteBeaker, TempBeaker
@ -22,8 +21,9 @@ def get_sha512(filename: str) -> str:
return hashlib.sha512(file.read()).hexdigest() return hashlib.sha512(file.read()).hexdigest()
@dataclass(frozen=True, eq=True) class Transform(BaseModel):
class Transform: model_config = ConfigDict(frozen=True)
name: str name: str
transform_func: Callable transform_func: Callable
error_map: dict[tuple, str] error_map: dict[tuple, str]
@ -58,7 +58,12 @@ class Recipe:
def __repr__(self) -> str: def __repr__(self) -> str:
return f"Recipe({self.name})" return f"Recipe({self.name})"
def add_beaker(self, name: str, datatype: type | None) -> Beaker: def add_beaker(
self,
name: str,
datatype: Type[BaseModel],
beaker_type: Type[Beaker] = SqliteBeaker,
) -> Beaker:
self.graph.add_node(name, datatype=datatype) self.graph.add_node(name, datatype=datatype)
if datatype is None: if datatype is None:
self.beakers[name] = TempBeaker(name, datatype, self) self.beakers[name] = TempBeaker(name, datatype, self)
@ -133,30 +138,30 @@ class Recipe:
for seed in seeds: for seed in seeds:
self.beakers[beaker_name].add_items(seed) self.beakers[beaker_name].add_items(seed)
# def get_metadata(self, table_name: str) -> dict: def get_metadata(self, table_name: str) -> dict:
# cursor = self.db.cursor() cursor = self.db.cursor()
# cursor.execute( cursor.execute(
# "SELECT data FROM _metadata WHERE table_name = ?", "SELECT data FROM _metadata WHERE table_name = ?",
# (table_name,), (table_name,),
# ) )
# try: try:
# data = cursor.fetchone()["data"] data = cursor.fetchone()["data"]
# log.debug("get_metadata", table_name=table_name, data=data) log.debug("get_metadata", table_name=table_name, data=data)
# return json.loads(data) return json.loads(data)
# except TypeError: except TypeError:
# log.debug("get_metadata", table_name=table_name, data={}) log.debug("get_metadata", table_name=table_name, data={})
# return {} return {}
# def save_metadata(self, table_name: str, data: dict) -> None: def save_metadata(self, table_name: str, data: dict) -> None:
# data_json = json.dumps(data) data_json = json.dumps(data)
# log.info("save_metadata", table_name=table_name, data=data_json) log.info("save_metadata", table_name=table_name, data=data_json)
# # sqlite upsert # sqlite upsert
# cursor = self.db.cursor() cursor = self.db.cursor()
# cursor.execute( cursor.execute(
# "INSERT INTO _metadata (table_name, data) VALUES (?, ?) ON CONFLICT(table_name) DO UPDATE SET data = ?", "INSERT INTO _metadata (table_name, data) VALUES (?, ?) ON CONFLICT(table_name) DO UPDATE SET data = ?",
# (table_name, data_json, data_json), (table_name, data_json, data_json),
# ) )
# self.db.commit() self.db.commit()
def csv_to_beaker(self, filename: str, beaker_name: str) -> None: def csv_to_beaker(self, filename: str, beaker_name: str) -> None:
beaker = self.beakers[beaker_name] beaker = self.beakers[beaker_name]