Merge pull request #36 from bryanforbes/improve-sql-dml

Improve `sql.dml`

Author: CaselIT

sqlalchemy-stubs/sql/dml.pyi

@@ -1,31 +1,33 @@
 from typing import Any
 from typing import Optional
+from typing import TypeVar
+from typing import Union
 
-from sqlalchemy.types import NullType as NullType
-from . import coercions as coercions
-from . import roles as roles
-from .base import ColumnCollection as ColumnCollection
-from .base import CompileState as CompileState
-from .base import DialectKWArgs as DialectKWArgs
-from .base import Executable as Executable
-from .base import HasCompileState as HasCompileState
-from .elements import BooleanClauseList as BooleanClauseList
-from .elements import ClauseElement as ClauseElement
-from .elements import Null as Null
-from .selectable import HasCTE as HasCTE
-from .selectable import HasPrefixes as HasPrefixes
-from .visitors import InternalTraversal as InternalTraversal
-from .. import exc as exc
-from .. import util as util
-from ..util import collections_abc as collections_abc
+from . import roles
+from .base import CompileState
+from .base import DialectKWArgs
+from .base import Executable
+from .base import HasCompileState
+from .base import ImmutableColumnCollection
+from .elements import BooleanClauseList
+from .elements import ClauseElement
+from .schema import Table
+from .selectable import HasCTE
+from .selectable import HasPrefixes
+
+_UB = TypeVar("_UB", bound=UpdateBase)
+_VB = TypeVar("_VB", bound=ValuesBase)
+_I = TypeVar("_I", bound=Insert)
+_DWB = TypeVar("_DWB", bound=DMLWhereBase)
+_UD = TypeVar("_UD", bound=Update)
 
 class DMLState(CompileState):
     isupdate: bool = ...
     isdelete: bool = ...
     isinsert: bool = ...
     def __init__(self, statement: Any, compiler: Any, **kw: Any) -> None: ...
     @property
-    def dml_table(self): ...
+    def dml_table(self) -> Table: ...
 
 class InsertDMLState(DMLState):
     isinsert: bool = ...
@@ -37,7 +39,7 @@ class UpdateDMLState(DMLState):
     isupdate: bool = ...
     include_table_with_column_exprs: bool = ...
     statement: Any = ...
-    is_multitable: Any = ...
+    is_multitable: bool = ...
     def __init__(self, statement: Any, compiler: Any, **kw: Any) -> None: ...
 
 class DeleteDMLState(DMLState):
@@ -57,27 +59,26 @@ class UpdateBase(
     __visit_name__: str = ...
     named_with_column: bool = ...
     is_dml: bool = ...
-    def params(self, *arg: Any, **kw: Any) -> None: ...
-    def with_dialect_options(self, **opt: Any) -> None: ...
-    def bind(self): ...
     bind: Any = ...
-    def returning(self, *cols: Any) -> None: ...
+    def params(self: _UB, *arg: Any, **kw: Any) -> _UB: ...
+    def with_dialect_options(self: _UB, **opt: Any) -> _UB: ...
+    def returning(self: _UB, *cols: Any) -> _UB: ...
     @property
-    def exported_columns(self): ...
+    def exported_columns(self) -> ImmutableColumnCollection[Any]: ...
     def with_hint(
-        self,
+        self: _UB,
         text: Any,
         selectable: Optional[Any] = ...,
         dialect_name: str = ...,
-    ) -> None: ...
+    ) -> _UB: ...
 
 class ValuesBase(UpdateBase):
     __visit_name__: str = ...
     select: Any = ...
     table: Any = ...
     def __init__(self, table: Any, values: Any, prefixes: Any) -> None: ...
-    def values(self, *args: Any, **kwargs: Any) -> None: ...
-    def return_defaults(self, *cols: Any) -> None: ...
+    def values(self: _VB, *args: Any, **kwargs: Any) -> _VB: ...
+    def return_defaults(self: _VB, *cols: Any) -> _VB: ...
 
 class Insert(ValuesBase):
     __visit_name__: str = ...
@@ -95,17 +96,19 @@ class Insert(ValuesBase):
         return_defaults: bool = ...,
         **dialect_kw: Any,
     ) -> None: ...
-    def inline(self) -> None: ...
+    def inline(self: _I) -> _I: ...
     def from_select(
-        self, names: Any, select: Any, include_defaults: bool = ...
-    ) -> None: ...
+        self: _I, names: Any, select: Any, include_defaults: bool = ...
+    ) -> _I: ...
 
 class DMLWhereBase:
-    def where(self, *whereclause: Any) -> None: ...
-    def filter(self, *criteria: Any): ...
-    def filter_by(self, **kwargs: Any): ...
+    def where(self: _DWB, *whereclause: Any) -> _DWB: ...
+    def filter(self: _DWB, *criteria: Any) -> _DWB: ...
+    def filter_by(self: _DWB, **kwargs: Any) -> _DWB: ...
     @property
-    def whereclause(self): ...
+    def whereclause(
+        self,
+    ) -> Optional[Union[BooleanClauseList[Any], ClauseElement]]: ...
 
 class Update(DMLWhereBase, ValuesBase):
     __visit_name__: str = ...
@@ -123,8 +126,8 @@ class Update(DMLWhereBase, ValuesBase):
         preserve_parameter_order: bool = ...,
         **dialect_kw: Any,
     ) -> None: ...
-    def ordered_values(self, *args: Any) -> None: ...
-    def inline(self) -> None: ...
+    def ordered_values(self: _UD, *args: Any) -> _UD: ...
+    def inline(self: _UD) -> _UD: ...
 
 class Delete(DMLWhereBase, UpdateBase):
     __visit_name__: str = ...