piccolo 1.9.0__py3-none-any.whl → 1.10.0__py3-none-any.whl

piccolo/engine/sqlite.py

@@ -11,7 +11,15 @@ from dataclasses import dataclass
 from decimal import Decimal
 from functools import partial, wraps
 
-from piccolo.engine.base import Batch, Engine, validate_savepoint_name
+from typing_extensions import Self
+
+from piccolo.engine.base import (
+    BaseAtomic,
+    BaseBatch,
+    BaseTransaction,
+    Engine,
+    validate_savepoint_name,
+)
 from piccolo.engine.exceptions import TransactionError
 from piccolo.query.base import DDL, Query
 from piccolo.querystring import QueryString
@@ -309,7 +317,7 @@ for column_name in ("TIMESTAMP", "TIMESTAMPTZ", "DATE", "TIME"):
 
 
 @dataclass
-class AsyncBatch(Batch):
+class AsyncBatch(BaseBatch):
     connection: Connection
     query: Query
     batch_size: int
@@ -327,16 +335,16 @@ class AsyncBatch(Batch):
         data = await self.cursor.fetchmany(self.batch_size)
         return await self.query._process_results(data)
 
-    def __aiter__(self):
+    def __aiter__(self: Self) -> Self:
         return self
 
-    async def __anext__(self):
+    async def __anext__(self) -> t.List[t.Dict]:
         response = await self.next()
         if response == []:
             raise StopAsyncIteration()
         return response
 
-    async def __aenter__(self):
+    async def __aenter__(self: Self) -> Self:
         querystring = self.query.querystrings[0]
         template, template_args = querystring.compile_string()
 
@@ -344,7 +352,7 @@ class AsyncBatch(Batch):
         return self
 
     async def __aexit__(self, exception_type, exception, traceback):
-        await self._cursor.close()
+        await self.cursor.close()
         await self.connection.close()
         return exception is not None
 
@@ -363,7 +371,7 @@ class TransactionType(enum.Enum):
     exclusive = "EXCLUSIVE"
 
 
-class Atomic:
+class Atomic(BaseAtomic):
     """
     Usage:
 
@@ -384,9 +392,9 @@ class Atomic:
     ):
         self.engine = engine
         self.transaction_type = transaction_type
-        self.queries: t.List[Query] = []
+        self.queries: t.List[t.Union[Query, DDL]] = []
 
-    def add(self, *query: Query):
+    def add(self, *query: t.Union[Query, DDL]):
         self.queries += list(query)
 
     async def run(self):
@@ -434,7 +442,7 @@ class Savepoint:
         )
 
 
-class SQLiteTransaction:
+class SQLiteTransaction(BaseTransaction):
     """
     Used for wrapping queries in a transaction, using a context manager.
     Currently it's async only.
@@ -534,7 +542,7 @@ class SQLiteTransaction:
 
     ###########################################################################
 
-    async def __aexit__(self, exception_type, exception, traceback):
+    async def __aexit__(self, exception_type, exception, traceback) -> bool:
         if self._parent:
             return exception is None
 
@@ -560,16 +568,8 @@ def dict_factory(cursor, row) -> t.Dict:
     return {col[0]: row[idx] for idx, col in enumerate(cursor.description)}
 
 
-class SQLiteEngine(Engine[t.Optional[SQLiteTransaction]]):
-    __slots__ = (
-        "connection_kwargs",
-        "current_transaction",
-        "log_queries",
-        "log_responses",
-    )
-
-    engine_type = "sqlite"
-    min_version_number = 3.25
+class SQLiteEngine(Engine[SQLiteTransaction]):
+    __slots__ = ("connection_kwargs",)
 
     def __init__(
         self,
@@ -613,7 +613,12 @@ class SQLiteEngine(Engine[t.Optional[SQLiteTransaction]]):
             f"sqlite_current_transaction_{path}", default=None
         )
 
-        super().__init__()
+        super().__init__(
+            engine_type="sqlite",
+            min_version_number=3.25,
+            log_queries=log_queries,
+            log_responses=log_responses,
+        )
 
     @property
     def path(self):

piccolo/query/methods/create_index.py

@@ -14,7 +14,7 @@ class CreateIndex(DDL):
     def __init__(
         self,
         table: t.Type[Table],
-        columns: t.List[t.Union[Column, str]],
+        columns: t.Union[t.List[Column], t.List[str]],
         method: IndexMethod = IndexMethod.btree,
         if_not_exists: bool = False,
         **kwargs,

piccolo/query/methods/drop_index.py

@@ -14,7 +14,7 @@ class DropIndex(Query):
     def __init__(
         self,
         table: t.Type[Table],
-        columns: t.List[t.Union[Column, str]],
+        columns: t.Union[t.List[Column], t.List[str]],
         if_exists: bool = True,
         **kwargs,
     ):

piccolo/query/methods/objects.py

@@ -5,7 +5,7 @@ import typing as t
 from piccolo.columns.column_types import ForeignKey
 from piccolo.columns.combination import And, Where
 from piccolo.custom_types import Combinable, TableInstance
-from piccolo.engine.base import Batch
+from piccolo.engine.base import BaseBatch
 from piccolo.query.base import Query
 from piccolo.query.methods.select import Select
 from piccolo.query.mixins import (
@@ -268,17 +268,17 @@ class Objects(
 
     ###########################################################################
 
-    def first(self: Self) -> First[TableInstance]:
+    def first(self) -> First[TableInstance]:
         self.limit_delegate.limit(1)
         return First[TableInstance](query=self)
 
-    def get(self: Self, where: Combinable) -> Get[TableInstance]:
+    def get(self, where: Combinable) -> Get[TableInstance]:
         self.where_delegate.where(where)
         self.limit_delegate.limit(1)
         return Get[TableInstance](query=First[TableInstance](query=self))
 
     def get_or_create(
-        self: Self,
+        self,
         where: Combinable,
         defaults: t.Optional[t.Dict[Column, t.Any]] = None,
     ) -> GetOrCreate[TableInstance]:
@@ -288,17 +288,17 @@ class Objects(
             query=self, table_class=self.table, where=where, defaults=defaults
         )
 
-    def create(self: Self, **columns: t.Any) -> Create[TableInstance]:
+    def create(self, **columns: t.Any) -> Create[TableInstance]:
         return Create[TableInstance](table_class=self.table, columns=columns)
 
     ###########################################################################
 
     async def batch(
-        self: Self,
+        self,
         batch_size: t.Optional[int] = None,
         node: t.Optional[str] = None,
         **kwargs,
-    ) -> Batch:
+    ) -> BaseBatch:
         if batch_size:
             kwargs.update(batch_size=batch_size)
         if node:

piccolo/query/methods/select.py

@@ -5,11 +5,11 @@ import typing as t
 from collections import OrderedDict
 
 from piccolo.columns import Column, Selectable
-from piccolo.columns.column_types import JSON, JSONB, PrimaryKey
+from piccolo.columns.column_types import JSON, JSONB
 from piccolo.columns.m2m import M2MSelect
 from piccolo.columns.readable import Readable
 from piccolo.custom_types import TableInstance
-from piccolo.engine.base import Batch
+from piccolo.engine.base import BaseBatch
 from piccolo.query.base import Query
 from piccolo.query.mixins import (
     AsOfDelegate,
@@ -217,7 +217,7 @@ class Select(Query[TableInstance, t.List[t.Dict[str, t.Any]]]):
         self,
         response: t.List[t.Dict[str, t.Any]],
         secondary_table: t.Type[Table],
-        secondary_table_pk: PrimaryKey,
+        secondary_table_pk: Column,
         m2m_name: str,
         m2m_select: M2MSelect,
         as_list: bool = False,
@@ -386,14 +386,20 @@ class Select(Query[TableInstance, t.List[t.Dict[str, t.Any]]]):
         return self
 
     @t.overload
-    def output(self: Self, *, as_list: bool) -> SelectList: ...
+    def output(self: Self, *, as_list: bool) -> SelectList:  # type: ignore
+        ...
 
     @t.overload
-    def output(self: Self, *, as_json: bool) -> SelectJSON: ...
+    def output(self: Self, *, as_json: bool) -> SelectJSON:  # type: ignore
+        ...
 
     @t.overload
     def output(self: Self, *, load_json: bool) -> Self: ...
 
+    @t.overload
+    def output(self: Self, *, load_json: bool, as_list: bool) -> SelectJSON:  # type: ignore  # noqa: E501
+        ...
+
     @t.overload
     def output(self: Self, *, nested: bool) -> Self: ...
 
@@ -404,7 +410,7 @@ class Select(Query[TableInstance, t.List[t.Dict[str, t.Any]]]):
         as_json: bool = False,
         load_json: bool = False,
         nested: bool = False,
-    ):
+    ) -> t.Union[Self, SelectJSON, SelectList]:
         self.output_delegate.output(
             as_list=as_list,
             as_json=as_json,
@@ -436,7 +442,7 @@ class Select(Query[TableInstance, t.List[t.Dict[str, t.Any]]]):
         batch_size: t.Optional[int] = None,
         node: t.Optional[str] = None,
         **kwargs,
-    ) -> Batch:
+    ) -> BaseBatch:
         if batch_size:
             kwargs.update(batch_size=batch_size)
         if node:

piccolo/query/mixins.py

@@ -207,7 +207,6 @@ class Returning:
 
 @dataclass
 class Output:
-
     as_json: bool = False
     as_list: bool = False
     as_objects: bool = False
@@ -236,7 +235,6 @@ class Callback:
 
 @dataclass
 class WhereDelegate:
-
     _where: t.Optional[Combinable] = None
     _where_columns: t.List[Column] = field(default_factory=list)
 
@@ -246,7 +244,8 @@ class WhereDelegate:
         needed.
         """
         self._where_columns = []
-        self._extract_columns(self._where)
+        if self._where is not None:
+            self._extract_columns(self._where)
         return self._where_columns
 
     def _extract_columns(self, combinable: Combinable):
@@ -277,7 +276,6 @@ class WhereDelegate:
 
 @dataclass
 class OrderByDelegate:
-
     _order_by: OrderBy = field(default_factory=OrderBy)
 
     def get_order_by_columns(self) -> t.List[Column]:
@@ -303,7 +301,6 @@ class OrderByDelegate:
 
 @dataclass
 class LimitDelegate:
-
     _limit: t.Optional[Limit] = None
     _first: bool = False
 
@@ -330,7 +327,6 @@ class AsOfDelegate:
 
 @dataclass
 class DistinctDelegate:
-
     _distinct: Distinct = field(
         default_factory=lambda: Distinct(enabled=False, on=None)
     )
@@ -356,7 +352,6 @@ class ReturningDelegate:
 
 @dataclass
 class CountDelegate:
-
     _count: bool = False
 
     def count(self):
@@ -365,7 +360,6 @@ class CountDelegate:
 
 @dataclass
 class AddDelegate:
-
     _add: t.List[Table] = field(default_factory=list)
 
     def add(self, *instances: Table, table_class: t.Type[Table]):
@@ -421,8 +415,7 @@ class OutputDelegate:
             self._output.nested = bool(nested)
 
     def copy(self) -> OutputDelegate:
-        _output = self._output.copy() if self._output is not None else None
-        return self.__class__(_output=_output)
+        return self.__class__(_output=self._output.copy())
 
 
 @dataclass

piccolo/schema.py

@@ -10,7 +10,6 @@ from piccolo.utils.sync import run_sync
 
 
 class SchemaDDLBase(abc.ABC):
-
     db: Engine
 
     @property
@@ -132,16 +131,19 @@ class ListTables:
         self.db = db
         self.schema_name = schema_name
 
-    async def run(self):
-        response = await self.db.run_querystring(
-            QueryString(
-                """
+    async def run(self) -> t.List[str]:
+        response = t.cast(
+            t.List[t.Dict],
+            await self.db.run_querystring(
+                QueryString(
+                    """
                 SELECT table_name
                 FROM information_schema.tables
                 WHERE table_schema = {}
                 """,
-                self.schema_name,
-            )
+                    self.schema_name,
+                )
+            ),
         )
         return [i["table_name"] for i in response]
 
@@ -156,9 +158,14 @@ class ListSchemas:
     def __init__(self, db: Engine):
         self.db = db
 
-    async def run(self):
-        response = await self.db.run_querystring(
-            QueryString("SELECT schema_name FROM information_schema.schemata")
+    async def run(self) -> t.List[str]:
+        response = t.cast(
+            t.List[t.Dict],
+            await self.db.run_querystring(
+                QueryString(
+                    "SELECT schema_name FROM information_schema.schemata"
+                )
+            ),
         )
         return [i["schema_name"] for i in response]
 
@@ -180,7 +187,7 @@ class SchemaManager:
         """
         db = db or engine_finder()
 
-        if not db:
+        if db is None:
             raise ValueError("The DB can't be found.")
 
         self.db = db

piccolo/table.py

@@ -143,8 +143,11 @@ class TableMeta:
     def db(self, value: Engine):
         self._db = value
 
-    def refresh_db(self):
-        self.db = engine_finder()
+    def refresh_db(self) -> None:
+        engine = engine_finder()
+        if engine is None:
+            raise ValueError("The engine can't be found")
+        self.db = engine
 
     def get_column_by_name(self, name: str) -> Column:
         """
@@ -184,8 +187,8 @@ class TableMeta:
 
 
 class TableMetaclass(type):
-    def __str__(cls):
-        return cls._table_str()
+    def __str__(cls) -> str:
+        return cls._table_str()  # type: ignore
 
     def __repr__(cls):
         """
@@ -822,7 +825,7 @@ class Table(metaclass=TableMetaclass):
     @classmethod
     def all_related(
         cls, exclude: t.Optional[t.List[t.Union[str, ForeignKey]]] = None
-    ) -> t.List[Column]:
+    ) -> t.List[ForeignKey]:
         """
         Used in conjunction with ``objects`` queries. Just as we can use
         ``all_related`` on a ``ForeignKey``, you can also use it for the table
@@ -1251,7 +1254,7 @@ class Table(metaclass=TableMetaclass):
     @classmethod
     def create_index(
         cls,
-        columns: t.List[t.Union[Column, str]],
+        columns: t.Union[t.List[Column], t.List[str]],
         method: IndexMethod = IndexMethod.btree,
         if_not_exists: bool = False,
     ) -> CreateIndex:
@@ -1273,7 +1276,9 @@ class Table(metaclass=TableMetaclass):
 
     @classmethod
     def drop_index(
-        cls, columns: t.List[t.Union[Column, str]], if_exists: bool = True
+        cls,
+        columns: t.Union[t.List[Column], t.List[str]],
+        if_exists: bool = True,
     ) -> DropIndex:
         """
         Drop a table index. If multiple columns are specified, this refers
@@ -1464,22 +1469,18 @@ async def drop_db_tables(*tables: t.Type[Table]) -> None:
         # SQLite doesn't support CASCADE, so we have to drop them in the
         # correct order.
         sorted_table_classes = reversed(sort_table_classes(list(tables)))
-        atomic = engine.atomic()
-        atomic.add(
-            *[
-                Alter(table=table).drop_table(if_exists=True)
-                for table in sorted_table_classes
-            ]
-        )
+        ddl_statements = [
+            Alter(table=table).drop_table(if_exists=True)
+            for table in sorted_table_classes
+        ]
     else:
-        atomic = engine.atomic()
-        atomic.add(
-            *[
-                table.alter().drop_table(cascade=True, if_exists=True)
-                for table in tables
-            ]
-        )
+        ddl_statements = [
+            table.alter().drop_table(cascade=True, if_exists=True)
+            for table in tables
+        ]
 
+    atomic = engine.atomic()
+    atomic.add(*ddl_statements)
     await atomic.run()
 
 

piccolo/utils/encoding.py

@@ -19,13 +19,15 @@ def dump_json(data: t.Any, pretty: bool = False) -> str:
             orjson_params["option"] = (
                 orjson.OPT_INDENT_2 | orjson.OPT_APPEND_NEWLINE  # type: ignore
             )
-        return orjson.dumps(data, **orjson_params).decode("utf8")
+        return orjson.dumps(data, **orjson_params).decode(  # type: ignore
+            "utf8"
+        )
     else:
         params: t.Dict[str, t.Any] = {"default": str}
         if pretty:
             params["indent"] = 2
-        return json.dumps(data, **params)
+        return json.dumps(data, **params)  # type: ignore
 
 
 def load_json(data: str) -> t.Any:
-    return orjson.loads(data) if ORJSON else json.loads(data)
+    return orjson.loads(data) if ORJSON else json.loads(data)  # type: ignore

{piccolo-1.9.0.dist-info → piccolo-1.10.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: piccolo
-Version: 1.9.0
+Version: 1.10.0
 Summary: A fast, user friendly ORM and query builder which supports asyncio.
 Home-page: https://github.com/piccolo-orm/piccolo
 Author: Daniel Townsend