From 4c15e1261c31a4676860e839b8f0b8b17d340334 Mon Sep 17 00:00:00 2001 From: Roland Walker Date: Tue, 25 Aug 2026 07:12:54 -0400 Subject: [PATCH] completions on Polars methods in transforms Add completions on Polars methods in .| dataframe transforms. Although completions are found dynamically, the user's code should not be evaluated, which means that completions may be limited in scope and correctness for complex expressions. Motivation: help Pandas users adjust to the slightly different Polars API. Followups could include completions on Altair plotting methods, which are not included here. --- changelog.md | 1 + mycli/packages/polars_completion.py | 293 +++++++++++++++++++++++++ mycli/packages/polars_transform.py | 60 +++-- mycli/sqlcompleter.py | 13 ++ test/pytests/test_polars_completion.py | 259 ++++++++++++++++++++++ 5 files changed, 602 insertions(+), 24 deletions(-) create mode 100644 mycli/packages/polars_completion.py create mode 100644 test/pytests/test_polars_completion.py diff --git a/changelog.md b/changelog.md index e3a5b7900..bca087019 100644 --- a/changelog.md +++ b/changelog.md @@ -5,6 +5,7 @@ Features -------- * Sort completion candidates by frecency from history. * Add help snippets in special /command completions. +* Add completions on Polars methods in dataframe transforms. Internal diff --git a/mycli/packages/polars_completion.py b/mycli/packages/polars_completion.py new file mode 100644 index 000000000..e0e199f0a --- /dev/null +++ b/mycli/packages/polars_completion.py @@ -0,0 +1,293 @@ +from __future__ import annotations + +import ast +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from functools import lru_cache +import inspect +from io import StringIO +import re +import sys +import tokenize +from types import ModuleType, UnionType +from typing import Any, Union, cast, get_args, get_origin + +import sqlglot + +from mycli.packages.polars_transform import _pipeline_operator_indexes + +_ATTRIBUTE_PATTERN = re.compile(r'(?s)(.*)\.([A-Za-z_][A-Za-z0-9_]*)?\s*$') +_ANNOTATION_NAME_PATTERN = re.compile(r'[A-Za-z_][A-Za-z0-9_]*') +_SENTINEL = '__mycli_polars_completion__' +_OPEN_TO_CLOSE = {'(': ')', '[': ']', '{': '}'} +_CLOSE = frozenset(_OPEN_TO_CLOSE.values()) + + +@dataclass(frozen=True, slots=True) +class PolarsCompletion: + text: str + display: str + display_meta: str + start_position: int + + +def complete_polars_transform(command: str) -> list[PolarsCompletion] | None: + """Return safe attribute completions for a Polars transform expression.""" + try: + tokens = sqlglot.tokenize(command) + pipe_index, output_index = _pipeline_operator_indexes(command, tokens, require_operands=False) + except sqlglot.errors.TokenError: + return None + if pipe_index is None: + return None + if output_index is not None: + return [] + + pipe = tokens[pipe_index + 1] + expression = command[pipe.end + 1 :].strip() + match = _ATTRIBUTE_PATTERN.fullmatch(expression) + if match is None: + return [] + prefix = match.group(2) or '' + source = f'{match.group(1)}.{_SENTINEL}' + closed_source = _close_open_delimiters(source) + if closed_source is None: + return [] + try: + parsed = ast.parse(closed_source, mode='eval') + except SyntaxError: + return [] + target = next( + (node for node in ast.walk(parsed) if isinstance(node, ast.Attribute) and node.attr == _SENTINEL), + None, + ) + if target is None: + return [] + + polars = _load_polars() + if polars is None: + return [] + owner = _infer_type(target.value, polars) + if owner is None: + return [] + return [ + PolarsCompletion( + text=f'{name}(' if is_callable else name, + display=name, + display_meta=metadata, + start_position=-len(prefix), + ) + for name, is_callable, metadata in _members(owner, polars) + if name.startswith(prefix) + ] + + +def _close_open_delimiters(source: str) -> str | None: + stack: list[str] = [] + try: + tokens = tokenize.generate_tokens(StringIO(source).readline) + for token in tokens: + if token.type != tokenize.OP: + continue + if token.string in _OPEN_TO_CLOSE: + stack.append(token.string) + elif token.string in _CLOSE: + if not stack or _OPEN_TO_CLOSE[stack.pop()] != token.string: + return None + except (IndentationError, tokenize.TokenError): + pass + return source + ''.join(_OPEN_TO_CLOSE[value] for value in reversed(stack)) + + +@lru_cache(maxsize=1) +def _load_polars() -> Any | None: + try: + import polars as pl + except ImportError: + return None + return pl + + +def _infer_type(node: ast.AST, polars: Any) -> Any | None: + if isinstance(node, ast.Name): + if node.id == 'df': + return polars.DataFrame + if node.id == 'pl': + return polars + return None + if isinstance(node, ast.Subscript): + owner = _infer_type(node.value, polars) + if owner is polars.DataFrame: + return polars.Series + if owner is polars.Expr: + return polars.Expr + return None + if isinstance(node, ast.Attribute): + owner = _infer_type(node.value, polars) + member = _static_member(owner, node.attr) + if member is None: + return None + if isinstance(member, property): + return _return_type(member.fget, owner, polars) + if inspect.isclass(member) or isinstance(member, ModuleType): + return member + return None + if isinstance(node, ast.Call): + if isinstance(node.func, ast.Attribute): + owner = _infer_type(node.func.value, polars) + member = _static_member(owner, node.func.attr) + return _call_return_type(member, owner, polars) + return None + if isinstance(node, (ast.BinOp, ast.BoolOp, ast.Compare, ast.UnaryOp)): + inferred = [_infer_type(child, polars) for child in ast.iter_child_nodes(node)] + if polars.Expr in inferred: + return polars.Expr + if polars.Series in inferred: + return polars.Series + return None + + +def _static_member(owner: Any | None, name: str) -> Any | None: + if owner is None: + return None + try: + member = inspect.getattr_static(owner, name) + except AttributeError: + return None + if isinstance(member, (classmethod, staticmethod)): + return member.__func__ + return member + + +def _call_return_type(member: Any | None, owner: Any, polars: Any) -> Any | None: + if member is None: + return None + if inspect.isclass(member): + return member + result = _return_type(member, owner, polars) + if result is not None: + return result + call = _static_member(type(member), '__call__') + return _return_type(call, owner, polars) + + +def _return_type(callable_object: Any | None, owner: Any, polars: Any) -> Any | None: + if callable_object is None: + return None + if inspect.isroutine(callable_object): + annotations = getattr(callable_object, '__annotations__', {}) + else: + try: + annotations = inspect.getattr_static(callable_object, '__annotations__') + except AttributeError: + annotations = {} + if not isinstance(annotations, Mapping): + return None + annotation = annotations.get('return') + if annotation is None or annotation is inspect.Signature.empty: + return None + if annotation is Any: + return None + if inspect.isclass(annotation): + return annotation if _is_polars_type(annotation) else None + arguments = get_args(annotation) + if arguments: + origin = cast(Any, get_origin(annotation)) + if origin not in (Union, UnionType): + return None + for argument in arguments: + resolved = _resolve_annotation(argument, callable_object, owner, polars) + if resolved is not None: + return resolved + return _resolve_annotation(annotation, callable_object, owner, polars) + + +def _resolve_annotation(annotation: Any, callable_object: Any, owner: Any, polars: Any) -> Any | None: + if inspect.isclass(annotation): + return annotation if _is_polars_type(annotation) else None + if not isinstance(annotation, str): + return None + if annotation in ('Self', 'typing.Self'): + return owner + if '[' in annotation: + return None + namespace = getattr(callable_object, '__globals__', {}) if inspect.isroutine(callable_object) else {} + for name in _ANNOTATION_NAME_PATTERN.findall(annotation): + candidate = namespace.get(name) + if not inspect.isclass(candidate): + candidate = _static_member(polars, name) + if not inspect.isclass(candidate): + candidate = _polars_types_by_name(polars).get(name) + if inspect.isclass(candidate) and _is_polars_type(candidate): + return candidate + return None + + +def _is_polars_type(value: type[Any]) -> bool: + return value.__module__ == 'polars' or value.__module__.startswith('polars.') + + +@lru_cache(maxsize=1) +def _polars_types_by_name(polars: Any) -> Mapping[str, type[Any]]: + """Index classes in loaded Polars modules without importing annotation paths.""" + candidates: dict[str, type[Any] | None] = {} + modules = [ + module + for module_name, module in tuple(sys.modules.items()) + if module is not None and (module is polars or module_name.startswith('polars.')) + ] + for module in modules: + for value in vars(module).values(): + if not inspect.isclass(value) or not _is_polars_type(value): + continue + name = value.__name__ + existing = candidates.get(name, value) + candidates[name] = value if existing is value else None + return {name: value for name, value in candidates.items() if value is not None} + + +@lru_cache(maxsize=None) +def _members(owner: Any, polars: Any) -> tuple[tuple[str, bool, str], ...]: + members: list[tuple[str, bool, str]] = [] + for name in dir(owner): + if name.startswith('_'): + continue + member = _static_member(owner, name) + if member is None: + continue + if isinstance(member, property): + members.append((name, False, _property_metadata(member, owner, polars))) + elif callable(member): + members.append((name, True, _signature_metadata(member))) + return tuple(members) + + +def _property_metadata(member: property, owner: Any, polars: Any) -> str: + result = _return_type(member.fget, owner, polars) + return result.__name__ if inspect.isclass(result) else 'property' + + +def _signature_metadata(member: Callable[..., Any]) -> str: + callable_object: Any = member + try: + signature = inspect.signature(callable_object) + except (TypeError, ValueError): + callable_object = _static_member(type(member), '__call__') + try: + signature = inspect.signature(callable_object) + except (TypeError, ValueError): + return '()' + parameters = [parameter for parameter in signature.parameters.values() if parameter.name not in ('self', 'cls')] + rendered: list[str] = [] + for parameter in parameters[:4]: + text = parameter.name + if parameter.kind == inspect.Parameter.VAR_POSITIONAL: + text = f'*{text}' + elif parameter.kind == inspect.Parameter.VAR_KEYWORD: + text = f'**{text}' + elif parameter.default is not inspect.Parameter.empty: + text = f'{text}=...' + rendered.append(text) + if len(parameters) > 4: + rendered.append('...') + return f"({', '.join(rendered)})" diff --git a/mycli/packages/polars_transform.py b/mycli/packages/polars_transform.py index b623c8d00..e6d3bb57e 100644 --- a/mycli/packages/polars_transform.py +++ b/mycli/packages/polars_transform.py @@ -66,30 +66,7 @@ def parse_polars_transform(command: str) -> PolarsPipeline | None: except sqlglot.errors.TokenError as exc: raise PolarsTransformError(f'Unable to parse Polars transform: {exc}') from exc - pipe_index: int | None = None - parquet_index: int | None = None - for index, token in enumerate(tokens[:-1]): - following = tokens[index + 1] - if token.token_type != sqlglot.TokenType.DOT: - continue - if token.start == 0 or not command[token.start - 1].isspace(): - continue - if following.token_type not in (sqlglot.TokenType.PIPE, sqlglot.TokenType.GT): - continue - if following.end + 1 >= len(command): - if following.token_type == sqlglot.TokenType.PIPE: - raise PolarsTransformError('Polars transforms require a Python expression.') - raise PolarsTransformError('File saves require a destination path.') - if not command[following.end + 1].isspace(): - continue - if following.token_type == sqlglot.TokenType.PIPE: - if pipe_index is not None: - raise PolarsTransformError('Polars transforms support only one ".|" operator.') - pipe_index = index - else: - if parquet_index is not None: - raise PolarsTransformError('File saves support only one ".>" operator.') - parquet_index = index + pipe_index, parquet_index = _pipeline_operator_indexes(command, tokens) if pipe_index is None and parquet_index is None: return None @@ -138,6 +115,41 @@ def parse_polars_transform(command: str) -> PolarsPipeline | None: return PolarsPipeline(sql=sql, expression=expression, output_path=output_path, output_mode=output_mode) +def _pipeline_operator_indexes( + command: str, + tokens: list[sqlglot.Token], + *, + require_operands: bool = True, +) -> tuple[int | None, int | None]: + """Locate transform and redirect operators in tokenized command text.""" + + pipe_index: int | None = None + parquet_index: int | None = None + for index, token in enumerate(tokens[:-1]): + following = tokens[index + 1] + if token.token_type != sqlglot.TokenType.DOT: + continue + if token.start == 0 or not command[token.start - 1].isspace(): + continue + if following.token_type not in (sqlglot.TokenType.PIPE, sqlglot.TokenType.GT): + continue + if following.end + 1 >= len(command) and require_operands: + if following.token_type == sqlglot.TokenType.PIPE: + raise PolarsTransformError('Polars transforms require a Python expression.') + raise PolarsTransformError('File saves require a destination path.') + if following.end + 1 < len(command) and not command[following.end + 1].isspace(): + continue + if following.token_type == sqlglot.TokenType.PIPE: + if pipe_index is not None: + raise PolarsTransformError('Polars transforms support only one ".|" operator.') + pipe_index = index + else: + if parquet_index is not None: + raise PolarsTransformError('File saves support only one ".>" operator.') + parquet_index = index + return pipe_index, parquet_index + + def _parse_output_path(path: str) -> str: if path[0] in ('\'', '"'): if len(path) < 2 or path[-1] != path[0]: diff --git a/mycli/sqlcompleter.py b/mycli/sqlcompleter.py index 52679fb62..81e3797bb 100644 --- a/mycli/sqlcompleter.py +++ b/mycli/sqlcompleter.py @@ -19,6 +19,7 @@ from mycli.compat import WIN from mycli.packages.completion_engine import is_inside_quotes, suggest_type from mycli.packages.filepaths import complete_path, parse_path, suggest_path +from mycli.packages.polars_completion import complete_polars_transform from mycli.packages.ptoolkit.history import frecency_score from mycli.packages.special import llm from mycli.packages.special.dsn_aliases import DsnAliases @@ -1439,6 +1440,18 @@ def get_completions( complete_event: CompleteEvent | None, smart_completion: bool | None = None, ) -> Iterable[Completion]: + polars_completions = complete_polars_transform(document.text_before_cursor) + if polars_completions is not None: + return ( + Completion( + candidate.text, + candidate.start_position, + display=candidate.display, + display_meta=candidate.display_meta, + ) + for candidate in polars_completions + ) + word_before_cursor = document.get_word_before_cursor(WORD=True) last_for_len = last_word(word_before_cursor, include="most_punctuations") text_for_len = last_for_len.lower() diff --git a/test/pytests/test_polars_completion.py b/test/pytests/test_polars_completion.py new file mode 100644 index 000000000..f5d0a497f --- /dev/null +++ b/test/pytests/test_polars_completion.py @@ -0,0 +1,259 @@ +from __future__ import annotations + +import ast +import builtins +from typing import Any + +from prompt_toolkit.completion import CompleteEvent +from prompt_toolkit.document import Document +import pytest +import sqlglot + +from mycli.packages import polars_completion +from mycli.packages.polars_completion import PolarsCompletion, complete_polars_transform +from mycli.sqlcompleter import SQLCompleter + + +def completion(command: str, text: str) -> PolarsCompletion: + candidates = complete_polars_transform(command) + assert candidates is not None + return next(candidate for candidate in candidates if candidate.text == text) + + +def test_completion_returns_none_outside_polars_transform() -> None: + assert complete_polars_transform("SELECT '.| df.fi'") is None + + +def test_completion_offers_dataframe_methods() -> None: + candidate = completion('SELECT * FROM orders .| df.fi', 'filter(') + + assert candidate.display == 'filter' + assert candidate.display_meta == '(*predicates, **constraints)' + assert candidate.start_position == -2 + + +def test_completion_offers_polars_functions_in_unclosed_call() -> None: + candidate = completion('SELECT * FROM orders .| df.filter(pl.c', 'col(') + + assert candidate.display == 'col' + assert candidate.display_meta == '(name, *more_names)' + assert candidate.start_position == -1 + + +def test_completion_infers_dataframe_method_return_type() -> None: + candidate = completion('SELECT * FROM orders .| df.filter(pl.col("id") > 1).he', 'head(') + + assert candidate.display_meta == '(n=...)' + + +@pytest.mark.parametrize( + 'expression', + ( + 'pl.col("name").str.to_', + 'df["name"].str.to_', + ), +) +def test_completion_infers_expression_and_series_namespaces(expression: str) -> None: + assert completion(f'SELECT * FROM orders .| {expression}', 'to_uppercase(').start_position == -3 + + +def test_completion_infers_lazyframe_methods() -> None: + assert completion('SELECT * FROM orders .| df.lazy().co', 'collect(').display == 'collect' + + +@pytest.mark.parametrize( + ('expression', 'candidate'), + ( + ('df.group_by_dynamic("timestamp", every="1d").a', 'agg('), + ('df.rolling("timestamp", period="1d").a', 'agg('), + ('pl.when(pl.col("value") > 0).then(1).o', 'otherwise('), + ('df.plot.b', 'bar('), + ('df["value"].plot.h', 'hist('), + ), +) +def test_completion_resolves_forward_declared_polars_types(expression: str, candidate: str) -> None: + assert completion(f'SELECT * FROM orders .| {expression}', candidate).display == candidate.removesuffix('(') + + +def test_completion_does_not_evaluate_expression(monkeypatch: pytest.MonkeyPatch) -> None: + def fail_eval(*args: Any, **kwargs: Any) -> Any: + raise AssertionError('User expression was evaluated.') + + monkeypatch.setattr(builtins, 'eval', fail_eval) + + candidate = completion('SELECT 1 .| pl.when(danger()).then(1).o', 'otherwise(') + assert candidate.display == 'otherwise' + + +def test_completion_does_not_treat_generic_return_as_element_type() -> None: + assert complete_polars_transform('SELECT * FROM orders .| df.get_columns().he') == [] + + +def test_completion_does_not_add_parenthesis_to_property() -> None: + candidate = completion('SELECT * FROM orders .| df.col', 'columns') + + assert candidate.display_meta == 'property' + + +def test_completion_excludes_private_members() -> None: + candidates = complete_polars_transform('SELECT * FROM orders .| df._') + + assert candidates == [] + + +@pytest.mark.parametrize( + 'command', + ( + 'SELECT * FROM orders .| df', + 'SELECT * FROM orders .| unknown.fi', + 'SELECT * FROM orders .| unknown().fi', + 'SELECT * FROM orders .| df.fi .> output.parquet', + 'SELECT * FROM orders .| df.fi .>', + ), +) +def test_completion_suppresses_sql_candidates_in_uncompletable_transform(command: str) -> None: + assert complete_polars_transform(command) == [] + + +def test_completion_handles_missing_optional_dependency(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(polars_completion, '_load_polars', lambda: None) + + assert complete_polars_transform('SELECT * FROM orders .| df.fi') == [] + + +def test_completion_handles_sqlglot_token_error(monkeypatch: pytest.MonkeyPatch) -> None: + def raise_token_error(command: str) -> list[sqlglot.Token]: + raise sqlglot.errors.TokenError(command) + + monkeypatch.setattr(polars_completion.sqlglot, 'tokenize', raise_token_error) + + assert complete_polars_transform('SELECT 1 .| df.fi') is None + + +@pytest.mark.parametrize('closed_source', (None, 'not valid Python!')) +def test_completion_handles_invalid_balanced_expression( + monkeypatch: pytest.MonkeyPatch, + closed_source: str | None, +) -> None: + monkeypatch.setattr(polars_completion, '_close_open_delimiters', lambda source: closed_source) + + assert complete_polars_transform('SELECT 1 .| df.fi') == [] + + +def test_completion_handles_missing_sentinel_node(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(polars_completion.ast, 'walk', lambda node: []) + + assert complete_polars_transform('SELECT 1 .| df.fi') == [] + + +def test_close_open_delimiters_rejects_mismatched_delimiters() -> None: + assert polars_completion._close_open_delimiters('df[(])') is None + + +def test_load_polars_returns_none_when_dependency_is_missing(monkeypatch: pytest.MonkeyPatch) -> None: + original_import = builtins.__import__ + + def missing_polars(name: str, *args: Any, **kwargs: Any) -> Any: + if name == 'polars': + raise ImportError + return original_import(name, *args, **kwargs) + + polars_completion._load_polars.cache_clear() + monkeypatch.setattr(builtins, '__import__', missing_polars) + assert polars_completion._load_polars() is None + polars_completion._load_polars.cache_clear() + + +@pytest.mark.parametrize( + ('expression', 'expected_name'), + ( + ('pl.col("name")[0]', 'Expr'), + ('pl[0]', None), + ('df.missing', None), + ('pl.DataFrame', 'DataFrame'), + ('df.filter', None), + ('unknown()', None), + ('pl.col("name") + 1', 'Expr'), + ('df["name"] + 1', 'Series'), + ('1 + 2', None), + ), +) +def test_infer_type_handles_static_expression_forms(expression: str, expected_name: str | None) -> None: + import polars as pl + + inferred = polars_completion._infer_type(ast.parse(expression, mode='eval').body, pl) + + assert getattr(inferred, '__name__', None) == expected_name + + +def test_static_member_handles_absent_owner_and_attribute() -> None: + assert polars_completion._static_member(None, 'value') is None + assert polars_completion._static_member(object, 'missing') is None + + +def test_call_return_type_handles_missing_member_and_class() -> None: + import polars as pl + + assert polars_completion._call_return_type(None, pl, pl) is None + assert polars_completion._call_return_type(pl.DataFrame, pl, pl) is pl.DataFrame + + +def test_return_type_handles_annotation_variants() -> None: + import polars as pl + + class InvalidAnnotations: + __annotations__ = 1 # type: ignore[assignment] + + def returns_any() -> Any: + raise AssertionError + + def returns_int() -> int: + raise AssertionError + + def returns_generic() -> list[pl.Series]: + raise AssertionError + + def returns_union() -> int | pl.DataFrame: + raise AssertionError + + returns_any.__annotations__['return'] = Any + returns_int.__annotations__['return'] = int + returns_generic.__annotations__['return'] = list[pl.Series] + returns_union.__annotations__['return'] = int | pl.DataFrame + + assert polars_completion._return_type(None, pl, pl) is None + assert polars_completion._return_type(InvalidAnnotations(), pl, pl) is None + assert polars_completion._return_type(returns_any, pl, pl) is None + assert polars_completion._return_type(returns_int, pl, pl) is None + assert polars_completion._return_type(returns_generic, pl, pl) is None + assert polars_completion._return_type(returns_union, pl, pl) is pl.DataFrame + + +def test_resolve_annotation_handles_direct_self_and_invalid_annotations() -> None: + import polars as pl + + assert polars_completion._resolve_annotation(pl.Expr, lambda: None, pl.DataFrame, pl) is pl.Expr + assert polars_completion._resolve_annotation(object(), lambda: None, pl.DataFrame, pl) is None + assert polars_completion._resolve_annotation('Self', lambda: None, pl.DataFrame, pl) is pl.DataFrame + + +def test_signature_metadata_handles_uninspectable_callable(monkeypatch: pytest.MonkeyPatch) -> None: + def fail_signature(member: Any) -> Any: + raise ValueError(member) + + monkeypatch.setattr(polars_completion.inspect, 'signature', fail_signature) + + assert polars_completion._signature_metadata(lambda: None) == '()' + + +@pytest.mark.parametrize('smart_completion', (True, False)) +def test_sqlcompleter_uses_polars_completion_in_all_modes(smart_completion: bool) -> None: + completer = SQLCompleter(smart_completion=smart_completion) + document = Document('SELECT * FROM orders .| df.fi') + + candidates = list(completer.get_completions(document, CompleteEvent())) + candidate = next(candidate for candidate in candidates if candidate.text == 'filter(') + + assert candidate.start_position == -2 + assert candidate.display_text == 'filter' + assert candidate.display_meta_text == '(*predicates, **constraints)'