Skip to content

Commit

Permalink
fix the return types of the classes' __enter__ functions (#384)
Browse files Browse the repository at this point in the history
fix the return types of the classes' __enter__ functions so that the type information is preserved in context managers eg with-as blocks

Signed-off-by: wyattscarpenter <[email protected]>
  • Loading branch information
wyattscarpenter authored Apr 2, 2024
1 parent 6e0fb78 commit 5636b8f
Showing 1 changed file with 4 additions and 2 deletions.
6 changes: 4 additions & 2 deletions src/databricks/sql/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -271,7 +271,8 @@ def _set_use_inline_params_with_warning(self, value: Union[bool, str]):

return value

def __enter__(self):
# The ideal return type for this method is perhaps Self, but that was not added until 3.11, and we support pre-3.11 pythons, currently.
def __enter__(self) -> "Connection":
return self

def __exit__(self, exc_type, exc_value, traceback):
Expand Down Expand Up @@ -409,7 +410,8 @@ def __init__(
self.escaper = ParamEscaper()
self.lastrowid = None

def __enter__(self):
# The ideal return type for this method is perhaps Self, but that was not added until 3.11, and we support pre-3.11 pythons, currently.
def __enter__(self) -> "Cursor":
return self

def __exit__(self, exc_type, exc_value, traceback):
Expand Down

0 comments on commit 5636b8f

Please sign in to comment.