getting close on types
This commit is contained in:
		
							parent
							
								
									d33483f136
								
							
						
					
					
						commit
						0998fe1e6a
					
				
					 3 changed files with 40 additions and 35 deletions
				
			
		|  | @ -12,7 +12,7 @@ PydanticModel = Type[BaseModel] | |||
| 
 | ||||
| 
 | ||||
| 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.model = model | ||||
|         self.recipe = recipe | ||||
|  | @ -42,7 +42,7 @@ class Beaker(abc.ABC): | |||
| 
 | ||||
| 
 | ||||
| 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) | ||||
|         self._items: list[tuple[str, BaseModel]] = [] | ||||
| 
 | ||||
|  | @ -62,11 +62,11 @@ class TempBeaker(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) | ||||
|         # create table if it doesn't exist | ||||
|         self.cursor = self.recipe.db.cursor() | ||||
|         self.cursor.row_factory = sqlite3.Row | ||||
|         self.cursor.row_factory = sqlite3.Row  # type: ignore | ||||
|         self.cursor.execute( | ||||
|             f"CREATE TABLE IF NOT EXISTS {self.name} (uuid TEXT PRIMARY KEY, data JSON)" | ||||
|         ) | ||||
|  |  | |||
|  | @ -67,7 +67,7 @@ def run( | |||
|     if not input and not has_data: | ||||
|         typer.secho("No data; pass --input to seed beaker(s)", fg=typer.colors.RED) | ||||
|         raise typer.Exit(1) | ||||
|     for input_str in input:  # type: ignore | ||||
|     for input_str in input or []: | ||||
|         beaker, filename = input_str.split("=") | ||||
|         ctx.obj.csv_to_beaker(filename, beaker) | ||||
|     ctx.obj.run_once(start, end) | ||||
|  |  | |||
|  | @ -5,11 +5,10 @@ import inspect | |||
| import sqlite3 | ||||
| import hashlib | ||||
| import asyncio | ||||
| import networkx | ||||
| import networkx  # type: ignore | ||||
| from collections import defaultdict, Counter | ||||
| from dataclasses import dataclass  # TODO: pydantic? | ||||
| from pydantic import BaseModel | ||||
| from typing import Iterable, Callable | ||||
| from typing import Iterable, Callable, Type | ||||
| from pydantic import BaseModel, ConfigDict | ||||
| from structlog import get_logger | ||||
| 
 | ||||
| from .beakers import Beaker, SqliteBeaker, TempBeaker | ||||
|  | @ -22,8 +21,9 @@ def get_sha512(filename: str) -> str: | |||
|         return hashlib.sha512(file.read()).hexdigest() | ||||
| 
 | ||||
| 
 | ||||
| @dataclass(frozen=True, eq=True) | ||||
| class Transform: | ||||
| class Transform(BaseModel): | ||||
|     model_config = ConfigDict(frozen=True) | ||||
| 
 | ||||
|     name: str | ||||
|     transform_func: Callable | ||||
|     error_map: dict[tuple, str] | ||||
|  | @ -58,7 +58,12 @@ class Recipe: | |||
|     def __repr__(self) -> str: | ||||
|         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) | ||||
|         if datatype is None: | ||||
|             self.beakers[name] = TempBeaker(name, datatype, self) | ||||
|  | @ -133,30 +138,30 @@ class Recipe: | |||
|             for seed in seeds: | ||||
|                 self.beakers[beaker_name].add_items(seed) | ||||
| 
 | ||||
|     # def get_metadata(self, table_name: str) -> dict: | ||||
|     #     cursor = self.db.cursor() | ||||
|     #     cursor.execute( | ||||
|     #         "SELECT data FROM _metadata WHERE table_name = ?", | ||||
|     #         (table_name,), | ||||
|     #     ) | ||||
|     #     try: | ||||
|     #         data = cursor.fetchone()["data"] | ||||
|     #         log.debug("get_metadata", table_name=table_name, data=data) | ||||
|     #         return json.loads(data) | ||||
|     #     except TypeError: | ||||
|     #         log.debug("get_metadata", table_name=table_name, data={}) | ||||
|     #         return {} | ||||
|     def get_metadata(self, table_name: str) -> dict: | ||||
|         cursor = self.db.cursor() | ||||
|         cursor.execute( | ||||
|             "SELECT data FROM _metadata WHERE table_name = ?", | ||||
|             (table_name,), | ||||
|         ) | ||||
|         try: | ||||
|             data = cursor.fetchone()["data"] | ||||
|             log.debug("get_metadata", table_name=table_name, data=data) | ||||
|             return json.loads(data) | ||||
|         except TypeError: | ||||
|             log.debug("get_metadata", table_name=table_name, data={}) | ||||
|             return {} | ||||
| 
 | ||||
|     # def save_metadata(self, table_name: str, data: dict) -> None: | ||||
|     #     data_json = json.dumps(data) | ||||
|     #     log.info("save_metadata", table_name=table_name, data=data_json) | ||||
|     #     # sqlite upsert | ||||
|     #     cursor = self.db.cursor() | ||||
|     #     cursor.execute( | ||||
|     #         "INSERT INTO _metadata (table_name, data) VALUES (?, ?) ON CONFLICT(table_name) DO UPDATE SET data = ?", | ||||
|     #         (table_name, data_json, data_json), | ||||
|     #     ) | ||||
|     #     self.db.commit() | ||||
|     def save_metadata(self, table_name: str, data: dict) -> None: | ||||
|         data_json = json.dumps(data) | ||||
|         log.info("save_metadata", table_name=table_name, data=data_json) | ||||
|         # sqlite upsert | ||||
|         cursor = self.db.cursor() | ||||
|         cursor.execute( | ||||
|             "INSERT INTO _metadata (table_name, data) VALUES (?, ?) ON CONFLICT(table_name) DO UPDATE SET data = ?", | ||||
|             (table_name, data_json, data_json), | ||||
|         ) | ||||
|         self.db.commit() | ||||
| 
 | ||||
|     def csv_to_beaker(self, filename: str, beaker_name: str) -> None: | ||||
|         beaker = self.beakers[beaker_name] | ||||
|  |  | |||
		Loading…
	
		Reference in a new issue
	
	 James Turk
						James Turk