From 6daaaa709b355514ac7e4dad56b0efc6c2dd1efe Mon Sep 17 00:00:00 2001 From: Marc Garcia Date: Thu, 5 Aug 2021 15:09:05 -0500 Subject: [PATCH 1/3] Fix mypy errors in backends/base --- ibis/backends/base/__init__.py | 4 ++-- ibis/backends/base/client.py | 8 ++++++-- ibis/backends/base/sql/alchemy/client.py | 18 +++++++++--------- ibis/backends/base/sql/alchemy/registry.py | 3 ++- ibis/backends/base/sql/client.py | 16 ++++++++++++++++ .../base/sql/compiler/query_builder.py | 5 ++++- ibis/backends/base/sql/compiler/translator.py | 3 ++- ibis/backends/base/sql/registry/window.py | 13 +++++++------ 8 files changed, 48 insertions(+), 22 deletions(-) diff --git a/ibis/backends/base/__init__.py b/ibis/backends/base/__init__.py index f610011ba0c1..75beb1f2a668 100644 --- a/ibis/backends/base/__init__.py +++ b/ibis/backends/base/__init__.py @@ -1,6 +1,6 @@ import abc import warnings -from typing import Any, Callable +from typing import Any, Callable, Type import ibis.expr.operations as ops import ibis.expr.types as ir @@ -20,7 +20,7 @@ class BaseBackend(abc.ABC): """ database_class = Database - table_class = ops.DatabaseTable + table_class: Type[ops.DatabaseTable] = ops.DatabaseTable @property @abc.abstractmethod diff --git a/ibis/backends/base/client.py b/ibis/backends/base/client.py index 52dd2fa56265..76698d7aa24c 100644 --- a/ibis/backends/base/client.py +++ b/ibis/backends/base/client.py @@ -1,3 +1,7 @@ +from __future__ import annotations + +from typing import List + import ibis.expr.types as ir @@ -19,7 +23,7 @@ def __repr__(self) -> str: """Return type name and the name of the database.""" return '{}({!r})'.format(type(self).__name__, self.name) - def __dir__(self) -> set: + def __dir__(self) -> List[str]: """Return a set of attributes and tables available for the database. Returns @@ -131,5 +135,5 @@ def list_tables(self, like: str = None) -> list: like=self._qualify_like(like), database=self.name ) - def _qualify_like(self, like: str) -> str: + def _qualify_like(self, like: str | None) -> str | None: return like diff --git a/ibis/backends/base/sql/alchemy/client.py b/ibis/backends/base/sql/alchemy/client.py index 7cda3b37fa52..6abd507704d1 100644 --- a/ibis/backends/base/sql/alchemy/client.py +++ b/ibis/backends/base/sql/alchemy/client.py @@ -1,6 +1,6 @@ import contextlib import warnings -from typing import List, Optional, Union +from typing import Dict, List, Optional, Union import pandas as pd import sqlalchemy as sa @@ -50,7 +50,7 @@ def __init__(self, con: sa.engine.Engine) -> None: self.con = con self.meta = sa.MetaData(bind=con) self._inspector = sa.inspect(con) - self._schemas = {} + self._schemas: Dict[str, sch.Schema] = {} @property def inspector(self): @@ -96,7 +96,7 @@ def begin(self): yield bind def create_table(self, name, expr=None, schema=None, database=None): - if database == self.database_name: + if database == self.current_database: # avoid fully qualified name database = None @@ -150,7 +150,7 @@ def drop_table( database: Optional[str] = None, force: bool = False, ) -> None: - if database == self.database_name: + if database == self.current_database: # avoid fully qualified name database = None @@ -200,7 +200,7 @@ def load_data( Loading data to a table from a different database is not yet implemented """ - if database == self.database_name: + if database == self.current_database: # avoid fully qualified name database = None @@ -214,7 +214,7 @@ def load_data( if self.has_attachment: # for database with attachment # see: https://github.com/ibis-project/ibis/issues/1930 - params['schema'] = self.database_name + params['schema'] = self.current_database data.to_sql( table_name, @@ -297,7 +297,7 @@ def list_tables( names = [x for x in names if like in x] return sorted(names) - def raw_sql(self, query: str): + def raw_sql(self, query: str, results=False): return _AutoCloseCursor(super().raw_sql(query)) def _log(self, sql): @@ -357,7 +357,7 @@ def insert( """ - if database == self.database_name: + if database == self.current_database: # avoid fully qualified name database = None @@ -371,7 +371,7 @@ def insert( if self.has_attachment: # for database with attachment # see: https://github.com/ibis-project/ibis/issues/1930 - params['schema'] = self.database_name + params['schema'] = self.current_database if isinstance(obj, pd.DataFrame): obj.to_sql( diff --git a/ibis/backends/base/sql/alchemy/registry.py b/ibis/backends/base/sql/alchemy/registry.py index 771ec45669f8..459f707ae6d9 100644 --- a/ibis/backends/base/sql/alchemy/registry.py +++ b/ibis/backends/base/sql/alchemy/registry.py @@ -1,4 +1,5 @@ import operator +from typing import Any, Dict import sqlalchemy as sa import sqlalchemy.sql as sql @@ -405,7 +406,7 @@ def _sort_key(t, expr): return sort_direction(t.translate(by)) -sqlalchemy_operation_registry = { +sqlalchemy_operation_registry: Dict[Any, Any] = { ops.And: fixed_arity(sql.and_, 2), ops.Or: fixed_arity(sql.or_, 2), ops.Not: unary(sa.not_), diff --git a/ibis/backends/base/sql/client.py b/ibis/backends/base/sql/client.py index bdf7c6a983c8..f76603cd0bd9 100644 --- a/ibis/backends/base/sql/client.py +++ b/ibis/backends/base/sql/client.py @@ -18,6 +18,22 @@ class SQLClient(Client, metaclass=abc.ABCMeta): table_class = ops.DatabaseTable table_expr_class = ir.TableExpr + @property + def con(self): + """ + Subclasses must set a `con` attribute with a DBAPI2 connection. + + This is defined here to make mypy happy. + """ + raise NotImplementedError( + 'Subclasses of SQLClient must set a `con` attribute with the ' + 'connection to the database' + ) + + @con.setter + def con(self, value): + self.con + def table(self, name, database=None): """Create a table expression. diff --git a/ibis/backends/base/sql/compiler/query_builder.py b/ibis/backends/base/sql/compiler/query_builder.py index c6ea44a35c74..657245930000 100644 --- a/ibis/backends/base/sql/compiler/query_builder.py +++ b/ibis/backends/base/sql/compiler/query_builder.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from io import StringIO import toolz @@ -8,6 +10,7 @@ import ibis.util as util from ibis.backends.base.sql.registry import quote_identifier from ibis.config import options +from ibis.expr.signaure import Argument from .base import DML, QueryAST, SetOp from .select_builder import SelectBuilder @@ -500,7 +503,7 @@ def _get_keyword_list(self): return ["EXCEPT"] * (len(self.tables) - 1) -def flatten_union(table: ir.TableExpr): +def flatten_union(table: ir.TableExpr | Argument): """Extract all union queries from `table`. Parameters diff --git a/ibis/backends/base/sql/compiler/translator.py b/ibis/backends/base/sql/compiler/translator.py index 3f11bfe71ec1..fb6f04c73a8b 100644 --- a/ibis/backends/base/sql/compiler/translator.py +++ b/ibis/backends/base/sql/compiler/translator.py @@ -1,4 +1,5 @@ import operator +from typing import Callable, Dict import ibis import ibis.common.exceptions as com @@ -201,7 +202,7 @@ class ExprTranslator: """ _registry = operation_registry - _rewrites = {} + _rewrites: Dict[ops.Node, Callable] = {} def __init__(self, expr, context, named=False, permit_subquery=False): self.expr = expr diff --git a/ibis/backends/base/sql/registry/window.py b/ibis/backends/base/sql/registry/window.py index 3b2de6fcda93..903b44937317 100644 --- a/ibis/backends/base/sql/registry/window.py +++ b/ibis/backends/base/sql/registry/window.py @@ -7,6 +7,7 @@ import ibis.expr.datatypes as dt import ibis.expr.operations as ops import ibis.expr.types as ir +from ibis.expr.signature import Argument _map_interval_to_microseconds = { 'W': 604800000000, @@ -43,7 +44,7 @@ def _replace_interval_with_scalar( expr: Union[ir.Expr, dt.Interval, float] -) -> Union[ir.FloatingScalar, float]: +) -> Union[ir.Expr, float, Argument]: """ Good old Depth-First Search to identify the Interval and IntervalValue components of the expression and return a comparable scalar expression. @@ -57,9 +58,9 @@ def _replace_interval_with_scalar( ------- preceding : float or ir.FloatingScalar, depending upon the expr """ - try: + if isinstance(expr, ir.Expr): expr_op = expr.op() - except AttributeError: + else: expr_op = None if not isinstance(expr, (dt.Interval, ir.IntervalValue)): @@ -80,13 +81,13 @@ def _replace_interval_with_scalar( ) elif expr_op.args and isinstance(expr, ir.IntervalValue): if len(expr_op.args) > 2: - raise com.NotImplementedError( - "'preceding' argument cannot be parsed." - ) + raise NotImplementedError("'preceding' argument cannot be parsed.") left_arg = _replace_interval_with_scalar(expr_op.args[0]) right_arg = _replace_interval_with_scalar(expr_op.args[1]) method = _map_interval_op_to_op[type(expr_op)] return method(left_arg, right_arg) + else: + raise TypeError(f'expr has unknown type {type(expr).__name__}') def cumulative_to_window(translator, expr, window): From 5b352bcf940ce83b89cfefe712b5870c454fad8f Mon Sep 17 00:00:00 2001 From: Marc Garcia Date: Thu, 5 Aug 2021 15:41:21 -0500 Subject: [PATCH 2/3] Fix typo --- ibis/backends/base/sql/compiler/query_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ibis/backends/base/sql/compiler/query_builder.py b/ibis/backends/base/sql/compiler/query_builder.py index 657245930000..000e6afdd123 100644 --- a/ibis/backends/base/sql/compiler/query_builder.py +++ b/ibis/backends/base/sql/compiler/query_builder.py @@ -10,7 +10,7 @@ import ibis.util as util from ibis.backends.base.sql.registry import quote_identifier from ibis.config import options -from ibis.expr.signaure import Argument +from ibis.expr.signature import Argument from .base import DML, QueryAST, SetOp from .select_builder import SelectBuilder From 9a110eae7487423e9c338bfbeee3d0dfaf6481e7 Mon Sep 17 00:00:00 2001 From: Marc Garcia Date: Fri, 6 Aug 2021 09:26:25 -0500 Subject: [PATCH 3/3] Fix tests and missing type validation --- ibis/backends/base/sql/client.py | 21 ++++--------------- .../base/sql/compiler/query_builder.py | 10 ++++++--- 2 files changed, 11 insertions(+), 20 deletions(-) diff --git a/ibis/backends/base/sql/client.py b/ibis/backends/base/sql/client.py index f76603cd0bd9..30dd5e59e096 100644 --- a/ibis/backends/base/sql/client.py +++ b/ibis/backends/base/sql/client.py @@ -18,22 +18,6 @@ class SQLClient(Client, metaclass=abc.ABCMeta): table_class = ops.DatabaseTable table_expr_class = ir.TableExpr - @property - def con(self): - """ - Subclasses must set a `con` attribute with a DBAPI2 connection. - - This is defined here to make mypy happy. - """ - raise NotImplementedError( - 'Subclasses of SQLClient must set a `con` attribute with the ' - 'connection to the database' - ) - - @con.setter - def con(self, value): - self.con - def table(self, name, database=None): """Create a table expression. @@ -118,7 +102,10 @@ def raw_sql(self, query: str, results=False): """ # TODO results is unused, it can be removed # (requires updating Impala tests) - cursor = self.con.execute(query) + # TODO `self.con` is assumed to be defined in subclasses, but there + # is nothing that enforces it. We should find a way to make sure + # `self.con` is always a DBAPI2 connection, or raise an error + cursor = self.con.execute(query) # type: ignore if cursor: return cursor cursor.release() diff --git a/ibis/backends/base/sql/compiler/query_builder.py b/ibis/backends/base/sql/compiler/query_builder.py index 000e6afdd123..5e904ce587c4 100644 --- a/ibis/backends/base/sql/compiler/query_builder.py +++ b/ibis/backends/base/sql/compiler/query_builder.py @@ -10,7 +10,6 @@ import ibis.util as util from ibis.backends.base.sql.registry import quote_identifier from ibis.config import options -from ibis.expr.signature import Argument from .base import DML, QueryAST, SetOp from .select_builder import SelectBuilder @@ -503,7 +502,7 @@ def _get_keyword_list(self): return ["EXCEPT"] * (len(self.tables) - 1) -def flatten_union(table: ir.TableExpr | Argument): +def flatten_union(table: ir.TableExpr): """Extract all union queries from `table`. Parameters @@ -516,8 +515,13 @@ def flatten_union(table: ir.TableExpr | Argument): """ op = table.op() if isinstance(op, ops.Union): + # For some reason mypy considers `op.left` and `op.right` + # of `Argument` type, and fails the validation. While in + # `flatten` types are the same, and it works return toolz.concatv( - flatten_union(op.left), [op.distinct], flatten_union(op.right) + flatten_union(op.left), # type: ignore + [op.distinct], + flatten_union(op.right), # type: ignore ) return [table]