From 2e20c86d06d971b73714f11e7e82476b292c3f26 Mon Sep 17 00:00:00 2001 From: wyattscarpenter Date: Tue, 2 Apr 2024 00:46:56 -0400 Subject: [PATCH] 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 --- src/databricks/sql/client.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index ad5146d1..e9e2978f 100644 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -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): @@ -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):