jupyter-duckdb 1.2.0.0__py3-none-any.whl → 1.4.111__py3-none-any.whl

duckdb_kernel/magics/MagicCommand.py

@@ -1,27 +1,28 @@
-from typing import Any, List, Tuple, Callable, Dict
+from typing import Any, List, Tuple, Callable, Dict, Set
 
 
 class MagicCommand:
-    _ARG = '''([^ ]+?|'.+?'|".+?")'''
+    _ARG = '''([^ ]+?|'.+?'|".+?")?'''
 
     def __init__(self, *names: str):
-        self._names: Tuple[str] = names
+        self._names: Tuple[str, ...] = names
 
-        self._arguments: List[Tuple[str, str]] = []
+        self._arguments: List[Tuple[str, Any, str]] = []
         self._flags: List[Tuple[str, str]] = []
         self._optionals: List[Tuple[str, Any, str]] = []
-
+        self._disables: Set[str] = set()
         self._code: bool = False
         self._result: bool = False
+        self._default: bool = False
 
         self._on: List[Callable] = []
 
     @property
-    def names(self) -> Tuple[str]:
+    def names(self) -> Tuple[str, ...]:
         return self._names
 
     @property
-    def args(self) -> List[Tuple[str, str]]:
+    def args(self) -> List[Tuple[str, Any, str]]:
         return self._arguments
 
     @property
@@ -32,6 +33,10 @@ class MagicCommand:
     def optionals(self) -> List[Tuple[str, Any, str]]:
         return self._optionals
 
+    @property
+    def disables(self) -> Set[str]:
+        return self._disables
+
     @property
     def requires_code(self) -> bool:
         return self._code
@@ -40,8 +45,17 @@ class MagicCommand:
     def requires_query_result(self) -> bool:
         return self._result
 
-    def arg(self, name: str, description: str = None) -> 'MagicCommand':
-        self._arguments.append((name, description))
+    @property
+    def is_default(self) -> bool:
+        return self._default
+
+    def arg(self, name: str, default_value: Any = None, description: str = None) -> 'MagicCommand':
+        if len(self._arguments) > 0:
+            ln, ldv, _ = self._arguments[-1]
+            if ldv is not None and default_value is None:
+                raise ValueError(f'argument {name} without default value registered after argument {ln} with default value {ldv}')
+
+        self._arguments.append((name, default_value, description))
         return self
 
     def opt(self, name: str, default_value: Any = None, description: str = None) -> 'MagicCommand':
@@ -52,6 +66,12 @@ class MagicCommand:
         self._flags.append((name, description))
         return self
 
+    def disable(self, *name: str) -> 'MagicCommand':
+        for n in name:
+            self._disables.add(n)
+
+        return self
+
     def code(self, code: bool) -> 'MagicCommand':
         self._code = code
         return self
@@ -60,10 +80,14 @@ class MagicCommand:
         self._result = result
         return self
 
-    def on(self, fun: Callable):
+    def on(self, fun: Callable) -> 'MagicCommand':
         self._on.append(fun)
         return self
 
+    def default(self, default: bool) -> 'MagicCommand':
+        self._default = default
+        return self
+
     @property
     def parameters(self) -> str:
         args = ' +'.join([self._ARG] * len(self._arguments))

duckdb_kernel/magics/MagicCommandCallback.py

@@ -1,20 +1,24 @@
 from typing import Optional, List
 
-from . import MagicCommand
+from . import MagicCommand, MagicState
 
 
 class MagicCommandCallback:
-    def __init__(self, mc: MagicCommand, silent: bool, code: str, *args, **kwargs):
+    def __init__(self, mc: MagicCommand, silent: bool, state: MagicState, *args, **kwargs):
         self._mc: MagicCommand = mc
         self._silent: bool = silent
-        self._code: str = code
+        self._state: MagicState = state
         self._args = args
         self._kwargs = kwargs
 
+    @property
+    def magic(self) -> MagicCommand:
+        return self._mc
+
     def __call__(self, columns: Optional[List[str]] = None, rows: Optional[List[List]] = None):
         if self._mc.requires_code:
-            return self._mc(self._silent, self._code, *self._args, **self._kwargs)
-        if self._mc.requires_query_result:
-            return self._mc(self._silent, columns, rows, *self._args, **self._kwargs)
+            self._mc(self._silent, self._state, *self._args, **self._kwargs)
+        elif self._mc.requires_query_result:
+            self._mc(self._silent, self._state, columns, rows, *self._args, **self._kwargs)
         else:
-            return self._mc(self._silent, *self._args, **self._kwargs)
+            self._mc(self._silent, self._state, *self._args, **self._kwargs)

duckdb_kernel/magics/MagicCommandHandler.py

@@ -2,6 +2,8 @@ import re
 from typing import Dict, Tuple, List
 
 from . import MagicCommand, MagicCommandException, MagicCommandCallback
+from .MagicState import MagicState
+from ..db import Connection
 
 
 class MagicCommandHandler:
@@ -14,20 +16,33 @@ class MagicCommandHandler:
                 key = key.lower()
                 self._magics[key] = cmd
 
-    def __call__(self, silent: bool, code: str) -> Tuple[str, List[MagicCommandCallback], List[MagicCommandCallback]]:
-        pre_query_callbacks = []
-        post_query_callbacks = []
+    def __getitem__(self, key: str) -> MagicCommand:
+        return self._magics[key.lower()]
+
+    def __call__(self, silent: bool, state: MagicState) \
+            -> Tuple[List[MagicCommandCallback], List[MagicCommandCallback]]:
+        enabled_callbacks: List[MagicCommandCallback] = []
+
+        # enable commands with default==True
+        for magic in self._magics.values():
+            if magic.is_default:
+                flags = {name: False for name, _ in magic.flags}
+                optionals = {name: default for name, default, _ in magic.optionals}
+                callback = MagicCommandCallback(magic, silent, state, **flags, **optionals)
 
+                enabled_callbacks.append(callback)
+
+        # search for magic commands in code
         while True:
             # ensure code starts with '%' or '%%' but not with '%%%'
-            match = re.match(r'^%{1,2}([^% ]+?)([ \t]*$| .+?$)', code, re.MULTILINE | re.IGNORECASE)
+            match = re.match(r'^%{1,2}([^% ]+?)([ \t]*$| .+?$)', state.code, re.MULTILINE | re.IGNORECASE)
 
             if match is None:
                 break
 
             # remove magic command from code
             start, end = match.span()
-            code = code[:start] + code[end + 1:]
+            state.code = state.code[:start] + state.code[end + 1:]
 
             # extract command
             command = match.group(1).lower()
@@ -45,7 +60,16 @@ class MagicCommandHandler:
                 raise MagicCommandException(f'could not parse parameters for command "{command}"')
 
             # extract args
-            args = [g for g, _ in zip(match.groups(), magic.args)]
+            args = [group if group is not None else default
+                    for group, (_, default, _) in zip(match.groups(), magic.args)]
+
+            args = [arg[1:-1]
+                    if arg is not None and (arg[0] == '"' and arg[-1] == '"' or arg[0] == "'" and arg[-1] == "'")
+                    else arg
+                    for arg in args]
+
+            if any(arg is None for arg in args):
+                raise MagicCommandException(f'could not parse parameters for command "{command}"')
 
             i = len(args) + 1
 
@@ -69,16 +93,41 @@ class MagicCommandHandler:
                 value = match.group(i + 2)
                 i += 3
 
+                if value is not None and (value[0] == '"' and value[-1] == '"' or value[0] == "'" and value[-1] == "'"):
+                    value = value[1:-1]
+
                 if name is not None:
                     optionals[name.lower()] = value
 
             # add to callbacks
-            callback = MagicCommandCallback(magic, silent, code, *args, **flags, **optionals)
+            callback = MagicCommandCallback(magic, silent, state, *args, **flags, **optionals)
+            enabled_callbacks.append(callback)
+
+        # disable overwritten callbacks
+        callbacks = []
+        blacklist = set()
+
+        for callback in reversed(enabled_callbacks):
+            for name in callback.magic.names:
+                if name in blacklist:
+                    break
+            else:
+                callbacks.append(callback)
+
+                for name in callback.magic.names:
+                    blacklist.add(name)
+                for disable in callback.magic.disables:
+                    blacklist.add(disable)
+
+        # prepare callback lists
+        pre_query_callbacks = []
+        post_query_callbacks = []
 
-            if not magic.requires_query_result:
+        for callback in reversed(callbacks):
+            if not callback.magic.requires_query_result:
                 pre_query_callbacks.append(callback)
             else:
                 post_query_callbacks.append(callback)
 
         # return callbacks
-        return code, pre_query_callbacks, post_query_callbacks
+        return pre_query_callbacks, post_query_callbacks

duckdb_kernel/magics/MagicState.py

@@ -0,0 +1,11 @@
+from typing import Union, Dict, Optional
+
+from ..db import Connection
+
+
+class MagicState:
+    def __init__(self, db: Connection, code: str, max_rows: Optional[int]):
+        self.db: Connection = db
+        self.code: Union[str, Dict] = code
+        self.max_rows: Optional[int] = max_rows
+        self.column_name_mapping: Dict[str, str] = {}

duckdb_kernel/magics/__init__.py

@@ -2,3 +2,4 @@ from .MagicCommand import MagicCommand
 from .MagicCommandCallback import MagicCommandCallback
 from .MagicCommandException import MagicCommandException
 from .MagicCommandHandler import MagicCommandHandler
+from .MagicState import MagicState

duckdb_kernel/parser/DCParser.py

@@ -1,14 +1,22 @@
+from .ParserError import DCParserError
 from .elements import *
 from .tokenizer import *
+from .util.QuerySplitter import get_last_query
 
 
 class DCParser:
     @staticmethod
     def parse_query(query: str) -> DC_SET:
+        # remove comments from query
+        query = get_last_query(query, split_at=None, remove_comments=True)
+
         # create initial token set
         initial_token = Token(query)
         tokens = tuple(Tokenizer.tokenize(initial_token))
 
+        if len(tokens) == 0:
+            return None
+
         # split at |
         for i, token in enumerate(tokens):
             if token in DC_SET.symbols():
@@ -18,17 +26,17 @@ class DCParser:
                 )
 
         # raise exception if query is not in the correct format
-        raise AssertionError('The expression shall be of the format "{ x1, ..., xn | f(x1, ..., xn) }".')
+        raise DCParserError('The expression shall be of the format "{ x1, ..., xn | f(x1, ..., xn) }".', 0)
 
     @staticmethod
-    def parse_projection(*tokens: Token) -> LogicOperand:
+    def parse_projection(*tokens: Token, depth: int = 0) -> LogicOperand:
         if len(tokens) == 1:
             tokens = tuple(Tokenizer.tokenize(tokens[0]))
 
         return LogicOperand(*tokens)
 
     @staticmethod
-    def parse_condition(*tokens: Token) -> LogicElement:
+    def parse_condition(*tokens: Token, depth: int = 0) -> LogicElement:
         if len(tokens) == 1:
             tokens = tuple(Tokenizer.tokenize(tokens[0]))
 
@@ -40,8 +48,8 @@ class DCParser:
                     # return the operator
                     # with left part of tokens and right part of tokens
                     return operator(
-                        DCParser.parse_condition(*tokens[:-i]),
-                        DCParser.parse_condition(*tokens[-i + 1:])
+                        DCParser.parse_condition(*tokens[:-i], depth=depth + 1),
+                        DCParser.parse_condition(*tokens[-i + 1:], depth=depth + 1)
                     )
 
         # not
@@ -56,10 +64,12 @@ class DCParser:
         elif len(tokens) == 2:
             return DCOperand(
                 tokens[0],
-                tuple(Tokenizer.tokenize(tokens[1]))
+                tuple(Tokenizer.tokenize(tokens[1])),
+                depth=depth + 1
             )
         else:
             return DCOperand(
                 tokens[0],
-                tokens[1:]
+                tokens[1:],
+                depth=depth + 1
             )

duckdb_kernel/parser/LogicParser.py

@@ -4,12 +4,12 @@ from .tokenizer import *
 
 class LogicParser:
     @staticmethod
-    def parse_query(query: str) -> LogicElement:
+    def parse_query(query: str, depth: int = 0) -> LogicElement:
         initial_token = Token(query)
-        return LogicParser.parse_tokens(initial_token)
+        return LogicParser.parse_tokens(initial_token, depth=depth)
 
     @staticmethod
-    def parse_tokens(*tokens: Token) -> LogicElement:
+    def parse_tokens(*tokens: Token, depth: int = 0) -> LogicElement:
         if len(tokens) == 1:
             tokens = tuple(Tokenizer.tokenize(tokens[0]))
 
@@ -21,14 +21,14 @@ class LogicParser:
                     # return the operator
                     # with left part of tokens and right part of tokens
                     return operator(
-                        LogicParser.parse_tokens(*tokens[:-i]),
-                        LogicParser.parse_tokens(*tokens[-i + 1:])
+                        LogicParser.parse_tokens(*tokens[:-i], depth=depth + 1),
+                        LogicParser.parse_tokens(*tokens[-i + 1:], depth=depth + 1)
                     )
 
         # not
         if tokens[0] in LOGIC_NOT.symbols():
             return LOGIC_NOT(
-                LogicParser.parse_tokens(*tokens[1:])
+                LogicParser.parse_tokens(*tokens[1:], depth=depth + 1)
             )
 
         # ArgList

duckdb_kernel/parser/ParserError.py

@@ -0,0 +1,18 @@
+class ParserError(Exception):
+    def __init__(self, message: str, depth: int):
+        super().__init__(message)
+
+        self.message: str = message
+        self.depth: int = depth
+
+
+class RAParserError(ParserError):
+    pass
+
+
+class DCParserError(ParserError):
+    pass
+
+
+class LogicParserError(ParserError):
+    pass

duckdb_kernel/parser/RAParser.py

@@ -1,6 +1,7 @@
-from .LogicParser import LogicParser
+from .ParserError import RAParserError
 from .elements import *
 from .tokenizer import *
+from .util.QuerySplitter import get_last_query
 
 
 # Instead of multiple nested loops, a tree with rotation can
@@ -8,31 +9,39 @@ from .tokenizer import *
 
 class RAParser:
     @staticmethod
-    def parse_query(query: str) -> RAElement:
+    def parse_query(query: str) -> RAElement | None:
+        # remove comments from query
+        query = get_last_query(query, split_at=None, remove_comments=True)
+
+        # parse query
         initial_token = Token(query)
-        return RAParser.parse_tokens(initial_token)
+        return RAParser.parse_tokens(initial_token, depth=0)
 
     @staticmethod
-    def parse_tokens(*tokens: Token, target: RAOperator | RAOperand = None) -> RAElement:
+    def parse_tokens(*tokens: Token, target: RAOperator | RAOperand = None, depth: int = 0) -> RAElement | None:
         if len(tokens) == 1:
             tokens = tuple(Tokenizer.tokenize(tokens[0]))
 
         # binary operators
-        for operator in RA_BINARY_OPERATORS:
+        for operator_symbols in RA_BINARY_SYMBOLS:
             # iterate tokens and match symbol
             for i in range(1, len(tokens) + 1):
-                if tokens[-i].lower() in operator.symbols():
+                lower_token = tokens[-i].lower()
+
+                if lower_token in operator_symbols:
+                    operator = operator_symbols[lower_token]
+
                     # raise error if left or right operand missing
                     if i == 1:
-                        raise AssertionError(f'right operand missing after {tokens[-i]}')
+                        raise RAParserError(f'right operand missing after {tokens[-i]}', depth)
                     if i == len(tokens):
-                        raise AssertionError(f'left operand missing before {tokens[-i]}')
+                        raise RAParserError(f'left operand missing before {tokens[-i]}', depth)
 
                     # return the operator
                     # with left part of tokens and right part of tokens
                     return operator(
-                        RAParser.parse_tokens(*tokens[:-i]),
-                        RAParser.parse_tokens(*tokens[-i + 1:])
+                        RAParser.parse_tokens(*tokens[:-i], depth=depth + 1),
+                        RAParser.parse_tokens(*tokens[-i + 1:], depth=depth + 1)
                     )
 
         # unary operators
@@ -44,8 +53,8 @@ class RAParser:
                     # the last token is the operators target.
                     if target is None:
                         op = operator(
-                            RAParser.parse_tokens(tokens[-1]),
-                            LogicParser.parse_tokens(*tokens[-i + 1:-1])
+                            RAParser.parse_tokens(tokens[-1], depth=depth + 1),
+                            operator.parse_args(*tokens[-i + 1:-1], depth=depth + 1)
                         )
 
                     # Otherwise the handed target is this operator's
@@ -53,23 +62,22 @@ class RAParser:
                     else:
                         op = operator(
                             target,
-                            LogicParser.parse_tokens(*tokens[-i + 1:])
+                            operator.parse_args(*tokens[-i + 1:], depth=depth + 1)
                         )
 
                     # If there are any more tokens the operator is
                     # the target for the next step.
                     if i < len(tokens):
-                        return RAParser.parse_tokens(
-                            *tokens[:-i],
-                            target=op
-                        )
+                        return RAParser.parse_tokens(*tokens[:-i], target=op, depth=depth + 1)
 
                     # Otherwise the operator is the return value.
                     else:
                         return op
 
         # return as name
-        if len(tokens) > 1:
-            raise AssertionError(f'{tokens=}')
-
-        return RAOperand(tokens[0])
+        if len(tokens) == 0:
+            return None
+        elif len(tokens) == 1:
+            return RAOperand(tokens[0])
+        else:
+            raise RAParserError(f'{tokens=}', depth)

duckdb_kernel/parser/__init__.py

@@ -1,3 +1,4 @@
 from .DCParser import DCParser
 from .LogicParser import LogicParser
 from .RAParser import RAParser
+from .ParserError import *

duckdb_kernel/parser/elements/DCOperand.py

@@ -1,13 +1,14 @@
 from typing import Tuple
 
 from .LogicOperand import LogicOperand
+from ..ParserError import DCParserError
 from ..tokenizer import Token
 
 
 class DCOperand(LogicOperand):
-    def __new__(cls, relation: Token, columns: Tuple[Token], skip_comma: bool = False):
+    def __new__(cls, relation: Token, columns: Tuple[Token, ...], skip_comma: bool = False, depth: int = 0):
         if not skip_comma and not all(t == ',' for i, t in enumerate(columns) if i % 2 == 1):
-            raise AssertionError('arguments must be separated by commas')
+            raise DCParserError('arguments must be separated by commas', 0)
 
         return tuple.__new__(
             cls,
@@ -18,9 +19,11 @@ class DCOperand(LogicOperand):
             ))
         )
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, relation: Token, columns: Tuple[Token, ...], skip_comma: bool = False, depth: int = 0):
         super().__init__()
-        self.invert = False
+
+        self.depth: int = depth
+        self.invert: bool = False
 
     @property
     def relation(self) -> Token:

duckdb_kernel/parser/elements/LogicElement.py

@@ -1,5 +1,3 @@
-from typing import Iterator
-
 from ..util.RenamableColumnList import RenamableColumnList
 
 

duckdb_kernel/parser/elements/RAElement.py

@@ -36,10 +36,13 @@ class RAElement:
         # if all columns are from the same relation we can skip the relation name
         if len(set(c.table for c in columns)) == 1:
             column_names = ', '.join(f'{c.current_name} AS "{c.name}"' for c in columns)
+            order_names = ', '.join(f'"{c.name}" ASC' for c in columns)
         else:
             column_names = ', '.join(f'{c.current_name} AS "{c.full_name}"' for c in columns)
+            order_names = ', '.join(f'"{c.full_name}" ASC' for c in columns)
 
-        return f'SELECT {column_names} FROM ({sql}) {self._name()}'
+        # create sql
+        return f'SELECT {column_names} FROM ({sql}) {self._name()} ORDER BY {order_names}'
 
     def to_sql_with_count(self, tables: Dict[str, Table]) -> str:
         sql, _ = self.to_sql(tables)

duckdb_kernel/parser/elements/RARelationReference.py

@@ -0,0 +1,86 @@
+import re
+
+from . import RAUnaryOperator
+from .LogicElement import LogicElement
+from ..ParserError import RAParserError
+from ..tokenizer import Token
+
+
+class RARelationReference(LogicElement):
+    @staticmethod
+    def parse_tokens(operator: type[RAUnaryOperator], *tokens: Token, depth: int = 0) -> 'RARelationReference':
+        try:
+            # If we get one single token, it should be like
+            # R -> "R"
+            # [ R ] -> "R"
+            # [ R(A, B, C) ] -> "R(A, B, C)"
+            # (A, B, C) -> "(A, B, C")
+            # [ (A, B, C) ] -> "(A, B, C)"
+            if len(tokens) == 1:
+                return RARelationReference._parse_one_token(*tokens)
+
+            # If we get two tokens, it should be like
+            # R(A, B, C) -> "R", "A, B, C"
+            # R A -> "R", "A"
+            # (The latter equals R(A), but we should think about rejecting this type.)
+            elif len(tokens) == 2:
+                return RARelationReference._parse_two_tokens(*tokens)
+
+            # Otherwise, the input is malformed.
+            else:
+                raise AssertionError()
+
+        except AssertionError:
+            raise RAParserError(f'malformed input for operator {operator.symbols()[0]} {tokens=}', depth=depth)
+
+    @staticmethod
+    def _parse_one_token(token: Token) -> 'RARelationReference':
+        match = re.fullmatch(r'^\s*([A-Za-z0-9]+)?\s*(\(?((\s*[A-Za-z0-9]+\s*,\s*)*(\s*[A-Za-z0-9]+\s*,?\s*))\)?)?\s*$', token)
+        if match is None:
+            raise AssertionError()
+
+        if match.group(1) is not None:
+            relation = match.group(1).strip()
+        else:
+            relation = None
+
+        if match.group(3) is not None:
+            attributes = [b for b in (a.strip() for a in match.group(3).split(',')) if b != '']
+        else:
+            attributes = None
+
+        if relation is None and attributes is None:
+            raise AssertionError()
+
+        return RARelationReference(relation, attributes)
+
+    @staticmethod
+    def _parse_two_tokens(token1: Token, token2: Token) -> 'RARelationReference':
+        # We expect the first token to be a relation name and the second one
+        # to be a list of column names separated by commas.
+        relation = token1.strip()
+        attributes = [b for b in (a.strip() for a in token2.split(',')) if b != '']
+
+        return RARelationReference(relation, attributes)
+
+    def __init__(self, relation: str | None, attributes: list[str] | None):
+        # check duplicated attributes
+        if attributes is not None:
+            for i in range(len(attributes)):
+                for k in range(i + 1, len(attributes)):
+                    if attributes[i] == attributes[k]:
+                        raise RAParserError(f'duplicate attribute {attributes[i]}', 0)
+                    if attributes[i].lower() == attributes[k].lower():
+                        raise RAParserError(f'duplicate attribute {attributes[i]}={attributes[k]}', 0)
+
+        # store
+        self.relation: str | None = relation
+        self.attributes: list[str] | None = attributes
+
+    def __str__(self) -> str:
+        if self.relation is not None and self.attributes is None:
+            return self.relation
+        elif self.relation is None and self.attributes is not None:
+            return f'({", ".join(self.attributes)})'
+        else:
+            return f'{self.relation}({", ".join(self.attributes)})'

duckdb_kernel/parser/elements/RAUnaryOperator.py

@@ -3,9 +3,15 @@ from typing import Iterator
 from .LogicElement import LogicElement
 from .RAElement import RAElement
 from .RAOperator import RAOperator
+from ..tokenizer import Token
 
 
 class RAUnaryOperator(RAOperator):
+    @classmethod
+    def parse_args(cls: type['RAUnaryOperator'], *tokens: Token, depth: int):
+        from .. import LogicParser
+        return LogicParser.parse_tokens(*tokens, depth=depth)
+
     def __init__(self, target: RAElement):
         self.target: RAElement = target
 

duckdb_kernel/parser/elements/__init__.py

@@ -9,8 +9,9 @@ from .RAOperator import RAOperator
 from .RABinaryOperator import RABinaryOperator
 from .RAUnaryOperator import RAUnaryOperator
 from .RAOperand import RAOperand
-from .binary import RA_BINARY_OPERATORS
+from .binary import RA_BINARY_OPERATORS, RA_BINARY_SYMBOLS
 from .unary import RA_UNARY_OPERATORS
+from .RARelationReference import RARelationReference
 
 from .DCOperand import DCOperand
 from .binary import DC_SET

duckdb_kernel/parser/elements/binary/And.py

@@ -8,7 +8,7 @@ class And(LogicOperator):
 
     @staticmethod
     def symbols() -> Tuple[str, ...]:
-        return '∧', 'and', 'and'
+        return '∧', 'and'
 
     @property
     def sql_symbol(self) -> str: