pytrilogy 0.0.3.6__tar.gz → 0.0.3.8__tar.gz

{pytrilogy-0.0.3.6/pytrilogy.egg-info → pytrilogy-0.0.3.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: pytrilogy
-Version: 0.0.3.6
+Version: 0.0.3.8
 Summary: Declarative, typed query language that compiles to SQL.
 Home-page: 
 Author: 

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.8/pytrilogy.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: pytrilogy
-Version: 0.0.3.6
+Version: 0.0.3.8
 Summary: Declarative, typed query language that compiles to SQL.
 Home-page: 
 Author: 

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.8}/pytrilogy.egg-info/SOURCES.txt

@@ -28,6 +28,7 @@ tests/test_select.py
 tests/test_show.py
 tests/test_statements.py
 tests/test_undefined_concept.py
+tests/test_user_functions.py
 tests/test_where_clause.py
 trilogy/__init__.py
 trilogy/compiler.py
@@ -103,6 +104,7 @@ trilogy/dialect/base.py
 trilogy/dialect/bigquery.py
 trilogy/dialect/common.py
 trilogy/dialect/config.py
+trilogy/dialect/dataframe.py
 trilogy/dialect/duckdb.py
 trilogy/dialect/enums.py
 trilogy/dialect/postgres.py

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.8}/tests/test_models.py

@@ -219,10 +219,14 @@ def test_comparison():
     # this should not error
     Comparison(left=1, right=[1, 2, 3], operator=ComparisonOperator.IN)
 
-    Comparison(left=1, right=ListWrapper([1, 2, 3]), operator=ComparisonOperator.IN)
+    Comparison(
+        left=1,
+        right=ListWrapper([1, 2, 3], type=DataType.INTEGER),
+        operator=ComparisonOperator.IN,
+    )
 
 
-def test_comparison():
+def test_comparison_invalid_type():
     try:
         Comparison(left=1, right="abc", operator=ComparisonOperator.EQ)
     except Exception as exc:

pytrilogy-0.0.3.8/tests/test_user_functions.py

@@ -0,0 +1,17 @@
+from trilogy import Dialects
+
+
+def test_user_function_def():
+    x = Dialects.DUCK_DB.default_executor()
+
+    results = x.execute_query(
+        """
+def percent_ratio(a, b, digits=3) -> round(a::float / b * 100, digits);
+
+                 
+select @percent_ratio(10, 100) as ratio;
+                 
+                 """
+    )
+
+    assert results.fetchall()[0].ratio == 10.0

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.8}/trilogy/__init__.py

@@ -4,6 +4,6 @@ from trilogy.dialect.enums import Dialects
 from trilogy.executor import Executor
 from trilogy.parser import parse
 
-__version__ = "0.0.3.6"
+__version__ = "0.0.3.8"
 
 __all__ = ["parse", "Executor", "Dialects", "Environment", "CONFIG"]

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.8}/trilogy/core/functions.py

@@ -503,32 +503,52 @@ FUNCTION_REGISTRY: dict[FunctionType, FunctionConfig] = {
         arg_count=1,
     ),
     FunctionType.ADD: FunctionConfig(
-        valid_inputs={DataType.INTEGER, DataType.FLOAT, DataType.NUMBER},
+        valid_inputs={
+            DataType.INTEGER,
+            DataType.FLOAT,
+            DataType.NUMBER,
+            DataType.NUMERIC,
+        },
         output_purpose=Purpose.PROPERTY,
         output_type=DataType.INTEGER,
         arg_count=InfiniteFunctionArgs,
     ),
     FunctionType.SUBTRACT: FunctionConfig(
-        valid_inputs={DataType.INTEGER, DataType.FLOAT, DataType.NUMBER},
+        valid_inputs={
+            DataType.INTEGER,
+            DataType.FLOAT,
+            DataType.NUMBER,
+            DataType.NUMERIC,
+        },
         output_purpose=Purpose.PROPERTY,
         output_type=DataType.INTEGER,
         arg_count=InfiniteFunctionArgs,
     ),
     FunctionType.MULTIPLY: FunctionConfig(
-        valid_inputs={DataType.INTEGER, DataType.FLOAT, DataType.NUMBER},
+        valid_inputs={
+            DataType.INTEGER,
+            DataType.FLOAT,
+            DataType.NUMBER,
+            DataType.NUMERIC,
+        },
         output_purpose=Purpose.PROPERTY,
         output_type=DataType.INTEGER,
         arg_count=InfiniteFunctionArgs,
     ),
     FunctionType.DIVIDE: FunctionConfig(
-        valid_inputs={DataType.INTEGER, DataType.FLOAT, DataType.NUMBER},
+        valid_inputs={
+            DataType.INTEGER,
+            DataType.FLOAT,
+            DataType.NUMBER,
+            DataType.NUMERIC,
+        },
         output_purpose=Purpose.PROPERTY,
         output_type=DataType.INTEGER,
         arg_count=InfiniteFunctionArgs,
     ),
     FunctionType.MOD: FunctionConfig(
         valid_inputs=[
-            {DataType.INTEGER, DataType.FLOAT, DataType.NUMBER},
+            {DataType.INTEGER, DataType.FLOAT, DataType.NUMBER, DataType.NUMERIC},
             {DataType.INTEGER},
         ],
         output_purpose=Purpose.PROPERTY,
@@ -537,7 +557,7 @@ FUNCTION_REGISTRY: dict[FunctionType, FunctionConfig] = {
     ),
     FunctionType.ROUND: FunctionConfig(
         valid_inputs=[
-            {DataType.INTEGER, DataType.FLOAT, DataType.NUMBER},
+            {DataType.INTEGER, DataType.FLOAT, DataType.NUMBER, DataType.NUMERIC},
             {DataType.INTEGER},
         ],
         output_purpose=Purpose.PROPERTY,

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.8}/trilogy/core/models/author.py

@@ -75,6 +75,9 @@ class Mergeable(ABC):
     def with_merge(self, source: Concept, target: Concept, modifiers: List[Modifier]):
         raise NotImplementedError
 
+    def with_reference_replacement(self, source: str, target: Expr):
+        raise NotImplementedError(type(self))
+
 
 class ConceptArgs(ABC):
     @property
@@ -152,6 +155,11 @@ class ConceptRef(Addressable, Namespaced, DataTyped, Mergeable, BaseModel):
             metadata=self.metadata,
         )
 
+    def with_reference_replacement(self, source: str, target: Expr):
+        if self.address == source:
+            return target
+        return self
+
 
 class UndefinedConcept(ConceptRef):
     pass
@@ -409,7 +417,7 @@ class Grain(Namespaced, BaseModel):
     ) -> Grain:
         from trilogy.parsing.common import concepts_to_grain_concepts
 
-        return Grain.model_construct(
+        x = Grain.model_construct(
             components={
                 c.address
                 for c in concepts_to_grain_concepts(concepts, environment=environment)
@@ -417,6 +425,8 @@ class Grain(Namespaced, BaseModel):
             where_clause=where_clause,
         )
 
+        return x
+
     def with_namespace(self, namespace: str) -> "Grain":
         return Grain.model_construct(
             components={address_with_namespace(c, namespace) for c in self.components},
@@ -662,6 +672,21 @@ class Comparison(ConceptArgs, Mergeable, Namespaced, BaseModel):
             operator=self.operator,
         )
 
+    def with_reference_replacement(self, source, target):
+        return self.__class__.model_construct(
+            left=(
+                self.left.with_reference_replacement(source, target)
+                if isinstance(self.left, Mergeable)
+                else self.left
+            ),
+            right=(
+                self.right.with_reference_replacement(source, target)
+                if isinstance(self.right, Mergeable)
+                else self.right
+            ),
+            operator=self.operator,
+        )
+
     def with_namespace(self, namespace: str):
         return self.__class__.model_construct(
             left=(
@@ -1233,19 +1258,9 @@ class OrderItem(Mergeable, ConceptArgs, Namespaced, BaseModel):
             order=self.order,
         )
 
-    @property
-    def output(self):
-        return self.expr.output
-
     @property
     def concept_arguments(self) -> Sequence[ConceptRef]:
-        base: List[ConceptRef] = []
-        x = self.expr
-        if isinstance(x, ConceptRef):
-            base += [x]
-        elif isinstance(x, ConceptArgs):
-            base += x.concept_arguments
-        return base
+        return get_concept_arguments(self.expr)
 
     @property
     def row_arguments(self) -> Sequence[ConceptRef]:
@@ -1316,31 +1331,19 @@ class WindowItem(DataTyped, ConceptArgs, Mergeable, Namespaced, BaseModel):
 
     @property
     def concept_arguments(self) -> List[ConceptRef]:
-        return self.arguments
-
-    @property
-    def arguments(self) -> List[ConceptRef]:
         output = [self.content]
         for order in self.order_by:
-            output += [order.output]
+            output += get_concept_arguments(order)
         for item in self.over:
-            output += [item]
+            output += get_concept_arguments(item)
         return output
 
-    @property
-    def output(self) -> ConceptRef:
-        return self.content
-
     @property
     def output_datatype(self):
         if self.type in (WindowType.RANK, WindowType.ROW_NUMBER):
             return DataType.INTEGER
         return self.content.output_datatype
 
-    @property
-    def output_purpose(self):
-        return Purpose.PROPERTY
-
 
 def get_basic_type(
     type: DataType | ListType | StructType | MapType | NumericType,
@@ -1407,12 +1410,28 @@ class CaseWhen(Namespaced, ConceptArgs, Mergeable, BaseModel):
             ),
         )
 
+    def with_reference_replacement(self, source, target):
+        return CaseWhen.model_construct(
+            comparison=self.comparison.with_reference_replacement(source, target),
+            expr=(
+                self.expr.with_reference_replacement(source, target)
+                if isinstance(self.expr, Mergeable)
+                else self.expr
+            ),
+        )
+
 
 class CaseElse(Namespaced, ConceptArgs, Mergeable, BaseModel):
     expr: "Expr"
     # this ensures that it's easily differentiable from CaseWhen
     discriminant: ComparisonOperator = ComparisonOperator.ELSE
 
+    def __str__(self):
+        return self.__repr__()
+
+    def __repr__(self):
+        return f"ELSE {str(self.expr)}"
+
     @field_validator("expr", mode="before")
     def enforce_expr(cls, v):
         if isinstance(v, Concept):
@@ -1435,6 +1454,19 @@ class CaseElse(Namespaced, ConceptArgs, Mergeable, BaseModel):
             ),
         )
 
+    def with_reference_replacement(self, source, target):
+        return CaseElse.model_construct(
+            discriminant=self.discriminant,
+            expr=(
+                self.expr.with_reference_replacement(
+                    source,
+                    target,
+                )
+                if isinstance(self.expr, Mergeable)
+                else self.expr
+            ),
+        )
+
     def with_namespace(self, namespace: str) -> CaseElse:
         return CaseElse.model_construct(
             discriminant=self.discriminant,
@@ -1510,8 +1542,6 @@ class Function(DataTyped, ConceptArgs, Mergeable, Namespaced, BaseModel):
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        if "datatype" in str(self):
-            raise SyntaxError(str(self))
 
     def __repr__(self):
         return f'{self.operator.value}({",".join([str(a) for a in self.arguments])})'
@@ -1597,6 +1627,29 @@ class Function(DataTyped, ConceptArgs, Mergeable, Namespaced, BaseModel):
                     )
         return v
 
+    def with_reference_replacement(self, source: str, target: Expr):
+        return Function.model_construct(
+            operator=self.operator,
+            arguments=[
+                (
+                    c.with_reference_replacement(
+                        source,
+                        target,
+                    )
+                    if isinstance(
+                        c,
+                        Mergeable,
+                    )
+                    else c
+                )
+                for c in self.arguments
+            ],
+            output_datatype=self.output_datatype,
+            output_purpose=self.output_purpose,
+            valid_inputs=self.valid_inputs,
+            arg_count=self.arg_count,
+        )
+
     def with_namespace(self, namespace: str) -> "Function":
         return Function.model_construct(
             operator=self.operator,
@@ -1658,10 +1711,13 @@ class Function(DataTyped, ConceptArgs, Mergeable, Namespaced, BaseModel):
         return base_grain
 
 
-class AggregateWrapper(Mergeable, ConceptArgs, Namespaced, BaseModel):
+class AggregateWrapper(Mergeable, DataTyped, ConceptArgs, Namespaced, BaseModel):
     function: Function
     by: List[ConceptRef] = Field(default_factory=list)
 
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
     @field_validator("by", mode="before")
     @classmethod
     def enforce_concept_ref(cls, v):
@@ -1693,10 +1749,6 @@ class AggregateWrapper(Mergeable, ConceptArgs, Namespaced, BaseModel):
     def output_purpose(self):
         return self.function.output_purpose
 
-    @property
-    def arguments(self):
-        return self.function.arguments
-
     def with_merge(self, source: Concept, target: Concept, modifiers: List[Modifier]):
         return AggregateWrapper.model_construct(
             function=self.function.with_merge(source, target, modifiers=modifiers),
@@ -1707,6 +1759,16 @@ class AggregateWrapper(Mergeable, ConceptArgs, Namespaced, BaseModel):
             ),
         )
 
+    def with_reference_replacement(self, source, target):
+        return AggregateWrapper.model_construct(
+            function=self.function.with_reference_replacement(source, target),
+            by=(
+                [c.with_reference_replacement(source, target) for c in self.by]
+                if self.by
+                else []
+            ),
+        )
+
     def with_namespace(self, namespace: str) -> "AggregateWrapper":
         return AggregateWrapper.model_construct(
             function=self.function.with_namespace(namespace),
@@ -1796,11 +1858,6 @@ class RowsetItem(Mergeable, ConceptArgs, Namespaced, BaseModel):
             rowset=self.rowset.with_namespace(namespace),
         )
 
-    @property
-    def arguments(self) -> List[ConceptRef]:
-        output = [self.content]
-        return output
-
     @property
     def output(self) -> ConceptRef:
         return self.content
@@ -1831,7 +1888,10 @@ class OrderBy(Mergeable, Namespaced, BaseModel):
 
     @property
     def concept_arguments(self):
-        return [x.expr for x in self.items]
+        base = []
+        for x in self.items:
+            base += x.concept_arguments
+        return base
 
 
 class AlignClause(Namespaced, BaseModel):
@@ -2088,6 +2148,11 @@ class Comment(BaseModel):
     text: str
 
 
+class ArgBinding(BaseModel):
+    name: str
+    default: Expr | None = None
+
+
 Expr = (
     MagicConstants
     | bool

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.8}/trilogy/core/models/build.py

@@ -1639,9 +1639,26 @@ class Factory:
 
     @build.register
     def _(self, base: Comparison) -> BuildComparison:
+        from trilogy.parsing.common import arbitrary_to_concept
+
+        left = base.left
+        if isinstance(left, AggregateWrapper):
+            left_c = arbitrary_to_concept(
+                left,
+                environment=self.environment,
+            )
+            left = left_c  # type: ignore
+        right = base.right
+        if isinstance(right, AggregateWrapper):
+            right_c = arbitrary_to_concept(
+                right,
+                environment=self.environment,
+            )
+
+            right = right_c  # type: ignore
         return BuildComparison.model_construct(
-            left=(self.build(base.left)),
-            right=(self.build(base.right)),
+            left=(self.build(left)),
+            right=(self.build(right)),
             operator=base.operator,
         )
 
@@ -1660,7 +1677,7 @@ class Factory:
         )
 
     @build.register
-    def _(self, base: RowsetItem):
+    def _(self, base: RowsetItem) -> BuildRowsetItem:
 
         factory = Factory(
             environment=self.environment,

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.8}/trilogy/core/models/datasource.py

@@ -117,6 +117,12 @@ class Datasource(HasUUID, Namespaced, BaseModel):
     where: Optional[WhereClause] = None
     non_partial_for: Optional[WhereClause] = None
 
+    @property
+    def safe_address(self) -> str:
+        if isinstance(self.address, Address):
+            return self.address.location
+        return self.address
+
     def __eq__(self, other):
         if not isinstance(other, Datasource):
             return False

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.8}/trilogy/core/models/environment.py

@@ -9,6 +9,7 @@ from typing import (
     TYPE_CHECKING,
     Annotated,
     Any,
+    Callable,
     Dict,
     ItemsView,
     List,
@@ -189,7 +190,7 @@ class Environment(BaseModel):
     datasources: Annotated[
         EnvironmentDatasourceDict, PlainValidator(validate_datasources)
     ] = Field(default_factory=EnvironmentDatasourceDict)
-    functions: Dict[str, Function] = Field(default_factory=dict)
+    functions: Dict[str, Callable[..., Any]] = Field(default_factory=dict)
     data_types: Dict[str, DataType] = Field(default_factory=dict)
     named_statements: Dict[str, SelectLineage] = Field(default_factory=dict)
     imports: Dict[str, list[Import]] = Field(

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.8}/trilogy/core/processing/node_generators/group_node.py

@@ -37,7 +37,7 @@ def gen_group_node(
         resolve_function_parent_concepts(concept, environment=environment), "address"
     )
     logger.info(
-        f"{padding(depth)}{LOGGER_PREFIX} parent concepts for {concept} are {[x.address for x in parent_concepts]} from group grain {concept.grain}"
+        f"{padding(depth)}{LOGGER_PREFIX} parent concepts for {concept} {concept.lineage} are {[x.address for x in parent_concepts]} from group grain {concept.grain}"
     )
 
     # if the aggregation has a grain, we need to ensure these are the ONLY optional in the output of the select

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.8}/trilogy/core/processing/node_generators/multiselect_node.py

@@ -76,7 +76,7 @@ def gen_multiselect_node(
     for select in lineage.selects:
 
         snode: StrategyNode = get_query_node(history.base_environment, select)
-        # raise SyntaxError(select.output_components)
+
         logger.info(
             f"{padding(depth)}{LOGGER_PREFIX} Fetched parent node with outputs {select.output_components}"
         )

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.8}/trilogy/core/statements/author.py

@@ -127,7 +127,9 @@ class SelectStatement(HasUUID, SelectTypeMixin, BaseModel):
             order_by=order_by,
             meta=meta or Metadata(),
         )
+
         output.grain = output.calculate_grain(environment)
+
         for x in selection:
 
             if x.is_undefined and environment.concepts.fail_on_missing:

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.8}/trilogy/dialect/config.py

@@ -1,3 +1,12 @@
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    try:
+        from pandas import DataFrame
+    except ImportError:
+        DataFrame = Any
+
+
 class DialectConfig:
     def __init__(self):
         pass
@@ -104,3 +113,9 @@ class TrinoConfig(PrestoConfig):
         if self.schema:
             return f"trino://{self.username}:{self.password}@{self.host}:{self.port}/{self.catalog}/{self.schema}"
         return f"trino://{self.username}:{self.password}@{self.host}:{self.port}/{self.catalog}"
+
+
+class DataFrameConfig(DuckDBConfig):
+    def __init__(self, dataframes: dict[str, "DataFrame"]):
+        super().__init__()
+        self.dataframes = dataframes

pytrilogy-0.0.3.8/trilogy/dialect/dataframe.py

@@ -0,0 +1,47 @@
+from typing import TYPE_CHECKING, Any
+
+from sqlalchemy import text
+
+from trilogy.core.models.environment import Environment
+from trilogy.dialect.duckdb import DuckDBDialect
+from trilogy.engine import ExecutionEngine
+
+if TYPE_CHECKING:
+    try:
+        from pandas import DataFrame
+    except ImportError:
+        DataFrame = Any
+
+
+class DataframeDialect(DuckDBDialect):
+    pass
+
+
+class DataframeConnectionWrapper(ExecutionEngine):
+    def __init__(self, engine: ExecutionEngine, dataframes: dict[str, "DataFrame"]):
+        self.engine = engine
+        self.dataframes = dataframes
+        self.connection = None
+
+    def setup(self, env: Environment, connection):
+        self._register_dataframes(env, connection)
+
+    def _register_dataframes(self, env: Environment, connection):
+        for ds in env.datasources.values():
+            if ds.safe_address in self.dataframes:
+                connection.execute(
+                    text("register(:name, :df)"),
+                    {"name": ds.safe_address, "df": self.dataframes[ds.safe_address]},
+                )
+            else:
+                raise ValueError(
+                    f"Dataframe {ds.safe_address} not found in dataframes on connection config, have {self.dataframes.keys()}"
+                )
+        pass
+
+    def add_dataframe(self, name: str, df: "DataFrame", connection, env: Environment):
+        self.dataframes[name] = df
+        self._register_dataframes(env, connection)
+
+    def connect(self) -> Any:
+        return self.engine.connect()

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.8}/trilogy/dialect/enums.py

@@ -16,7 +16,7 @@ def default_factory(conf: DialectConfig, config_type):
 
     if not isinstance(conf, config_type):
         raise TypeError(
-            f"Invalid dialect configuration for type {type(config_type).__name__}"
+            f"Invalid dialect configuration for type {type(config_type).__name__}, is {type(conf)}"
         )
     if conf.connect_args:
         return create_engine(
@@ -33,6 +33,7 @@ class Dialects(Enum):
     TRINO = "trino"
     POSTGRES = "postgres"
     SNOWFLAKE = "snowflake"
+    DATAFRAME = "dataframe"
 
     @classmethod
     def _missing_(cls, value):
@@ -88,6 +89,16 @@ class Dialects(Enum):
             from trilogy.dialect.config import TrinoConfig
 
             return _engine_factory(conf, TrinoConfig)
+        elif self == Dialects.DATAFRAME:
+            from trilogy.dialect.config import DataFrameConfig
+            from trilogy.dialect.dataframe import DataframeConnectionWrapper
+
+            if not conf:
+                conf = DataFrameConfig(dataframes={})
+
+            base = _engine_factory(conf, DataFrameConfig)
+
+            return DataframeConnectionWrapper(base, dataframes=conf.dataframes)
         else:
             raise ValueError(
                 f"Unsupported dialect {self} for default engine creation; create one explicitly."

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.8}/trilogy/engine.py

@@ -1,7 +1,9 @@
-from typing import Protocol
+from typing import Any, Protocol
 
 from sqlalchemy.engine import Connection, CursorResult, Engine
 
+from trilogy.core.models.environment import Environment
+
 
 class EngineResult(Protocol):
     pass
@@ -13,7 +15,7 @@ class EngineResult(Protocol):
 class EngineConnection(Protocol):
     pass
 
-    def execute(self, statement: str) -> EngineResult:
+    def execute(self, statement: str, parameters: Any | None = None) -> EngineResult:
         pass
 
 
@@ -23,6 +25,9 @@ class ExecutionEngine(Protocol):
     def connect(self) -> EngineConnection:
         pass
 
+    def setup(self, env: Environment, connection):
+        pass
+
 
 ### Begin default SQLAlchemy implementation
 class SqlAlchemyResult(EngineResult):
@@ -37,8 +42,10 @@ class SqlAlchemyConnection(EngineConnection):
     def __init__(self, connection: Connection):
         self.connection = connection
 
-    def execute(self, statement: str) -> SqlAlchemyResult:
-        return SqlAlchemyResult(self.connection.execute(statement))
+    def execute(
+        self, statement: str, parameters: Any | None = None
+    ) -> SqlAlchemyResult:
+        return SqlAlchemyResult(self.connection.execute(statement, parameters))
 
 
 class SqlAlchemyEngine(ExecutionEngine):