vtlengine 1.0__py3-none-any.whl → 1.0.1__py3-none-any.whl

vtlengine/Operators/Conditional.py

@@ -1,31 +1,37 @@
-import os
 from copy import copy
-from typing import Union
+from typing import Union, Any
 
 import numpy as np
 
-from vtlengine.DataTypes import Boolean, COMP_NAME_MAPPING, binary_implicit_promotion, \
-    SCALAR_TYPES_CLASS_REVERSE, Null
+from vtlengine.DataTypes import (
+    Boolean,
+    COMP_NAME_MAPPING,
+    binary_implicit_promotion,
+    SCALAR_TYPES_CLASS_REVERSE,
+    Null,
+)
 from vtlengine.Operators import Operator, Binary
 
 from vtlengine.Exceptions import SemanticError
 from vtlengine.Model import Scalar, DataComponent, Dataset, Role
 
-if os.environ.get("SPARK", False):
-    import pyspark.pandas as pd
-else:
-    import pandas as pd
+# if os.environ.get("SPARK", False):
+#     import pyspark.pandas as pd
+# else:
+#     import pandas as pd
+import pandas as pd
 
 
 class If(Operator):
     """
     If class:
-        `If-then-else <https://sdmx.org/wp-content/uploads/VTL-2.1-Reference-Manual.pdf#page=225&zoom=100,72,142>`_ operator
+        `If-then-else <https://sdmx.org/wp-content/uploads/VTL-2.1-Reference-Manual.pdf#page=225&zoom=100,72,142>`_ operator # noqa E501
         inherits from Operator, a superclass that contains general validate and evaluate class methods.
         It has the following class methods:
     Class methods:
-        evaluate: Evaluates if the operation is well constructed, checking the actual condition and dropping a boolean
-        result. The result will depend on the data class, such as datacomponent and dataset.
+        evaluate: Evaluates if the operation is well constructed, checking the actual condition and
+        dropping a boolean result.
+        The result will depend on the data class, such as datacomponent and dataset.
 
         component_level_evaluation: Returns a pandas dataframe with data to set the condition
 
@@ -37,79 +43,97 @@ class If(Operator):
     """
 
     @classmethod
-    def evaluate(cls, condition, true_branch, false_branch):
+    def evaluate(cls, condition: Any, true_branch: Any, false_branch: Any) -> Any:
         result = cls.validate(condition, true_branch, false_branch)
-        if isinstance(condition, DataComponent):
-            result.data = cls.component_level_evaluation(condition, true_branch, false_branch)
-        if isinstance(condition, Dataset):
-            result = cls.dataset_level_evaluation(result, condition, true_branch, false_branch)
+        if not isinstance(result, Scalar):
+            if isinstance(condition, DataComponent):
+                result.data = cls.component_level_evaluation(condition, true_branch, false_branch)
+            if isinstance(condition, Dataset):
+                result = cls.dataset_level_evaluation(result, condition, true_branch, false_branch)
         return result
 
     @classmethod
-    def component_level_evaluation(cls, condition: DataComponent,
-                                   true_branch: Union[DataComponent, Scalar],
-                                   false_branch: Union[DataComponent, Scalar]):
-        if isinstance(true_branch, Scalar):
-            true_data = pd.Series(true_branch.value, index=condition.data.index)
-        else:
-            true_data = true_branch.data.reindex(condition.data.index)
-        if isinstance(false_branch, Scalar):
-            false_data = pd.Series(false_branch.value, index=condition.data.index)
-        else:
-            false_data = false_branch.data.reindex(condition.data.index)
-
-        result = np.where(condition.data, true_data, false_data)
+    def component_level_evaluation(
+        cls, condition: DataComponent, true_branch: Any, false_branch: Any
+    ) -> Any:
+        result = None
+        if condition.data is not None:
+            if isinstance(true_branch, Scalar):
+                true_data = pd.Series(true_branch.value, index=condition.data.index)
+            else:
+                true_data = true_branch.data.reindex(condition.data.index)
+            if isinstance(false_branch, Scalar):
+                false_data = pd.Series(false_branch.value, index=condition.data.index)
+            else:
+                false_data = false_branch.data.reindex(condition.data.index)
+            result = np.where(condition.data, true_data, false_data)
         return pd.Series(result, index=condition.data.index)
 
     @classmethod
-    def dataset_level_evaluation(cls, result, condition, true_branch, false_branch):
+    def dataset_level_evaluation(
+        cls, result: Any, condition: Any, true_branch: Any, false_branch: Any
+    ) -> Dataset:
         ids = condition.get_identifiers_names()
         condition_measure = condition.get_measures_names()[0]
         true_data = condition.data[condition.data[condition_measure] == True]
         false_data = condition.data[condition.data[condition_measure] != True].fillna(False)
 
         if isinstance(true_branch, Dataset):
-            if len(true_data) > 0:
-                true_data = pd.merge(true_data, true_branch.data, on=ids, how='right',
-                                     suffixes=('_condition', ''))
+            if len(true_data) > 0 and true_branch.data is not None:
+                true_data = pd.merge(
+                    true_data, true_branch.data, on=ids, how="right", suffixes=("_condition", "")
+                )
             else:
                 true_data = pd.DataFrame(columns=true_branch.get_components_names())
         else:
             true_data[condition_measure] = true_data[condition_measure].apply(
-                lambda x: true_branch.value)
+                lambda x: true_branch.value
+            )
         if isinstance(false_branch, Dataset):
-            if len(false_data) > 0:
-                false_data = pd.merge(false_data, false_branch.data, on=ids, how='right',
-                                      suffixes=('_condition', ''))
+            if len(false_data) > 0 and false_branch.data is not None:
+                false_data = pd.merge(
+                    false_data, false_branch.data, on=ids, how="right", suffixes=("_condition", "")
+                )
             else:
                 false_data = pd.DataFrame(columns=false_branch.get_components_names())
         else:
             false_data[condition_measure] = false_data[condition_measure].apply(
-                lambda x: false_branch.value)
+                lambda x: false_branch.value
+            )
 
-        result.data = pd.concat([true_data, false_data],
-                                ignore_index=True).drop_duplicates().sort_values(by=ids)
+        result.data = (
+            pd.concat([true_data, false_data], ignore_index=True)
+            .drop_duplicates()
+            .sort_values(by=ids)
+        )
         if isinstance(result, Dataset):
-            drop_columns = [column for column in result.data.columns if
-                            column not in result.components.keys()]
+            drop_columns = [
+                column for column in result.data.columns if column not in result.components.keys()
+            ]
             result.data = result.data.dropna(subset=drop_columns).drop(columns=drop_columns)
         if isinstance(true_branch, Scalar) and isinstance(false_branch, Scalar):
             result.get_measures()[0].data_type = true_branch.data_type
             result.get_measures()[0].name = COMP_NAME_MAPPING[true_branch.data_type]
-            result.data = result.data.rename(
-                columns={condition_measure: result.get_measures()[0].name})
+            if result.data is not None:
+                result.data = result.data.rename(
+                    columns={condition_measure: result.get_measures()[0].name}
+                )
         return result
 
     @classmethod
-    def validate(cls, condition, true_branch, false_branch) -> Scalar | DataComponent | Dataset:
+    def validate(  # noqa: C901
+        cls, condition: Any, true_branch: Any, false_branch: Any
+    ) -> Union[Scalar, DataComponent, Dataset]:
         nullable = False
         left = true_branch
         right = false_branch
         if true_branch.__class__ != false_branch.__class__:
-            if (isinstance(true_branch, DataComponent) and isinstance(false_branch, Dataset)) or \
-                    (isinstance(true_branch, Dataset) and isinstance(false_branch, DataComponent)):
+            if (isinstance(true_branch, DataComponent) and isinstance(false_branch, Dataset)) or (
+                isinstance(true_branch, Dataset) and isinstance(false_branch, DataComponent)
+            ):
                 raise ValueError(
-                    "If then and else operands cannot be dataset and component respectively")
+                    "If then and else operands cannot be dataset and component respectively"
+                )
             if isinstance(true_branch, Scalar):
                 left = false_branch
                 right = true_branch
@@ -117,42 +141,48 @@ class If(Operator):
         # Datacomponent
         if isinstance(condition, DataComponent):
             if not condition.data_type == Boolean:
-                raise SemanticError("1-1-9-11", op=cls.op,
-                                    type=SCALAR_TYPES_CLASS_REVERSE[condition.data_type])
+                raise SemanticError(
+                    "1-1-9-11", op=cls.op, type=SCALAR_TYPES_CLASS_REVERSE[condition.data_type]
+                )
             if not isinstance(left, Scalar) or not isinstance(right, Scalar):
                 nullable = condition.nullable
             else:
-                if isinstance(left, Scalar) and left.data_type == Null:
-                    nullable = True
-                if isinstance(right, Scalar) and right.data_type == Null:
+                if left.data_type == Null or right.data_type == Null:
                     nullable = True
             if isinstance(left, DataComponent):
                 nullable |= left.nullable
             if isinstance(right, DataComponent):
                 nullable |= right.nullable
-            return DataComponent(name='result', data=None,
-                                 data_type=binary_implicit_promotion(left.data_type,
-                                                                     right.data_type),
-                                 role=Role.MEASURE, nullable=nullable)
+            return DataComponent(
+                name="result",
+                data=None,
+                data_type=binary_implicit_promotion(left.data_type, right.data_type),
+                role=Role.MEASURE,
+                nullable=nullable,
+            )
 
         # Dataset
         if isinstance(left, Scalar) and isinstance(right, Scalar):
-            raise SemanticError("1-1-9-12", op=cls.op, then_symbol=left.name,
-                                else_symbol=right.name)
+            raise SemanticError(
+                "1-1-9-12", op=cls.op, then_symbol=left.name, else_symbol=right.name
+            )
         if isinstance(left, DataComponent):
-            raise SemanticError("1-1-9-12", op=cls.op, then_symbol=left.name,
-                                else_symbol=right.name)
+            raise SemanticError(
+                "1-1-9-12", op=cls.op, then_symbol=left.name, else_symbol=right.name
+            )
         if isinstance(left, Scalar):
-            left.data_type = right.data_type = binary_implicit_promotion(left.data_type,
-                                                                         right.data_type)
-            return Dataset(name='result', components=copy(condition.components), data=None)
+            left.data_type = right.data_type = binary_implicit_promotion(
+                left.data_type, right.data_type
+            )
+            return Dataset(name="result", components=copy(condition.components), data=None)
         if left.get_identifiers() != condition.get_identifiers():
             raise SemanticError("1-1-9-10", op=cls.op, clause=left.name)
         if isinstance(right, Scalar):
             for component in left.get_measures():
                 if component.data_type != right.data_type:
-                    component.data_type = binary_implicit_promotion(component.data_type,
-                                                                    right.data_type)
+                    component.data_type = binary_implicit_promotion(
+                        component.data_type, right.data_type
+                    )
         if isinstance(right, Dataset):
             if left.get_identifiers() != condition.get_identifiers():
                 raise SemanticError("1-1-9-10", op=cls.op, clause=right.name)
@@ -160,80 +190,101 @@ class If(Operator):
                 raise SemanticError("1-1-9-13", op=cls.op, then=left.name, else_clause=right.name)
             for component in left.get_measures():
                 if component.data_type != right.components[component.name].data_type:
-                    component.data_type = right.components[component.name].data_type = \
-                        binary_implicit_promotion(component.data_type,
-                                                  right.components[component.name].data_type)
+                    component.data_type = right.components[component.name].data_type = (
+                        binary_implicit_promotion(
+                            component.data_type, right.components[component.name].data_type
+                        )
+                    )
         if isinstance(condition, Dataset):
             if len(condition.get_measures()) != 1:
                 raise SemanticError("1-1-9-4", op=cls.op, name=condition.name)
             if condition.get_measures()[0].data_type != Boolean:
-                raise SemanticError("1-1-9-5", op=cls.op, type=SCALAR_TYPES_CLASS_REVERSE[
-                    condition.get_measures()[0].data_type])
+                raise SemanticError(
+                    "1-1-9-5",
+                    op=cls.op,
+                    type=SCALAR_TYPES_CLASS_REVERSE[condition.get_measures()[0].data_type],
+                )
             if left.get_identifiers() != condition.get_identifiers():
                 raise SemanticError("1-1-9-6", op=cls.op)
         result_components = {comp_name: copy(comp) for comp_name, comp in left.components.items()}
-        return Dataset(name='result', components=result_components, data=None)
+        return Dataset(name="result", components=result_components, data=None)
 
 
 class Nvl(Binary):
     """
     Null class:
-        `Nvl <https://sdmx.org/wp-content/uploads/VTL-2.1-Reference-Manual.pdf#page=229&zoom=100,72,370>`_operator class.
+        `Nvl <https://sdmx.org/wp-content/uploads/VTL-2.1-Reference-Manual.pdf#page=229&zoom=100,72,370>`_operator class. # noqa E501
         It has the following class methods:
 
     Class methods:
-        Validate: Class method that validates if the operation at scalar, datacomponent or dataset level can be performed.
+        Validate: Class method that validates if the operation at scalar,
+        datacomponent or dataset level can be performed.
         Evaluate: Evaluates the actual operation, returning the result.
     """
 
     @classmethod
-    def evaluate(cls, left, right):
+    def evaluate(cls, left: Any, right: Any) -> Union[Scalar, DataComponent, Dataset]:
         result = cls.validate(left, right)
 
-        if isinstance(left, Scalar):
-            if pd.isnull(left):
+        if isinstance(left, Scalar) and isinstance(result, Scalar):
+            if pd.isnull(left):  # type: ignore[call-overload]
                 result.value = right.value
             else:
                 result.value = left.value
-            return result
-        if isinstance(right, Scalar):
-            result.data = left.data.fillna(right.value)
-        if isinstance(right, Dataset) or isinstance(right, DataComponent):
-            result.data = left.data.fillna(right.data)
-        if isinstance(result, Dataset):
-            result.data = result.data[result.get_components_names()]
+        else:
+            if not isinstance(result, Scalar):
+                if isinstance(right, Scalar):
+                    result.data = left.data.fillna(right.value)
+                else:
+                    result.data = left.data.fillna(right.data)
+                if isinstance(result, Dataset):
+                    result.data = result.data[result.get_components_names()]
         return result
 
     @classmethod
-    def validate(cls, left, right) -> Scalar | DataComponent | Dataset:
+    def validate(cls, left: Any, right: Any) -> Union[Scalar, DataComponent, Dataset]:
+        result_components = {}
         if isinstance(left, Scalar):
             if not isinstance(right, Scalar):
                 raise ValueError(
-                    "Nvl operation at scalar level must have scalar types on right (applicable) side")
+                    "Nvl operation at scalar level must have scalar "
+                    "types on right (applicable) side"
+                )
             cls.type_validation(left.data_type, right.data_type)
-            return Scalar(name='result', value=None, data_type=left.data_type)
+            return Scalar(name="result", value=None, data_type=left.data_type)
         if isinstance(left, DataComponent):
             if isinstance(right, Dataset):
                 raise ValueError(
-                    "Nvl operation at component level cannot have dataset type on right (applicable) side")
+                    "Nvl operation at component level cannot have "
+                    "dataset type on right (applicable) side"
+                )
             cls.type_validation(left.data_type, right.data_type)
-            return DataComponent(name='result', data=pd.Series(dtype=object),
-                                 data_type=left.data_type,
-                                 role=Role.MEASURE, nullable=False)
+            return DataComponent(
+                name="result",
+                data=pd.Series(dtype=object),
+                data_type=left.data_type,
+                role=Role.MEASURE,
+                nullable=False,
+            )
         if isinstance(left, Dataset):
             if isinstance(right, DataComponent):
                 raise ValueError(
-                    "Nvl operation at dataset level cannot have component type on right (applicable) side")
+                    "Nvl operation at dataset level cannot have component "
+                    "type on right (applicable) side"
+                )
             if isinstance(right, Scalar):
                 for component in left.get_measures():
                     cls.type_validation(component.data_type, right.data_type)
             if isinstance(right, Dataset):
                 for component in left.get_measures():
-                    cls.type_validation(component.data_type,
-                                        right.components[component.name].data_type)
-            result_components = {comp_name: copy(comp) for comp_name, comp in
-                                 left.components.items()
-                                 if comp.role != Role.ATTRIBUTE}
+                    cls.type_validation(
+                        component.data_type, right.components[component.name].data_type
+                    )
+            result_components = {
+                comp_name: copy(comp)
+                for comp_name, comp in left.components.items()
+                if comp.role != Role.ATTRIBUTE
+            }
             for comp in result_components.values():
                 comp.nullable = False
-            return Dataset(name='result', components=result_components, data=None)
+        return Dataset(name="result", components=result_components, data=None)

vtlengine/Operators/General.py

@@ -1,4 +1,4 @@
-from typing import Dict, List
+from typing import Dict, List, Any, Union
 
 import pandas as pd
 import sqlite3
@@ -20,38 +20,47 @@ class Membership(Binary):
     """
 
     @classmethod
-    def validate(cls, left_operand: Dataset, right_operand: str):
+    def validate(cls, left_operand: Any, right_operand: Any) -> Dataset:
         if right_operand not in left_operand.components:
-            raise SemanticError("1-1-1-10", op=cls.op, comp_name=right_operand,
-                                dataset_name=left_operand.name)
+            raise SemanticError(
+                "1-1-1-10", op=cls.op, comp_name=right_operand, dataset_name=left_operand.name
+            )
 
         component = left_operand.components[right_operand]
         if component.role in (Role.IDENTIFIER, Role.ATTRIBUTE):
             right_operand = COMP_NAME_MAPPING[component.data_type]
-            left_operand.components[right_operand] = Component(name=right_operand,
-                                                               data_type=component.data_type,
-                                                               role=Role.MEASURE,
-                                                               nullable=component.nullable)
+            left_operand.components[right_operand] = Component(
+                name=right_operand,
+                data_type=component.data_type,
+                role=Role.MEASURE,
+                nullable=component.nullable,
+            )
             if left_operand.data is not None:
                 left_operand.data[right_operand] = left_operand.data[component.name]
             left_operand.data[right_operand] = left_operand.data[component.name]
-
-        result_components = {name: comp for name, comp in left_operand.components.items()
-                             if comp.role == Role.IDENTIFIER or comp.name == right_operand}
+        result_components = {
+            name: comp
+            for name, comp in left_operand.components.items()
+            if comp.role == Role.IDENTIFIER or comp.name == right_operand
+        }
         result_dataset = Dataset(name="result", components=result_components, data=None)
         return result_dataset
 
     @classmethod
-    def evaluate(cls, left_operand: Dataset, right_operand: str,
-                 is_from_component_assignment=False) -> Dataset:
+    def evaluate(
+        cls, left_operand: Dataset, right_operand: str, is_from_component_assignment: bool = False
+    ) -> Union[DataComponent, Dataset]:
         result_dataset = cls.validate(left_operand, right_operand)
-        if is_from_component_assignment:
-            return DataComponent(name=right_operand,
-                                 data_type=left_operand.components[right_operand].data_type,
-                                 role=Role.MEASURE,
-                                 nullable=left_operand.components[right_operand].nullable,
-                                 data=left_operand.data[right_operand])
-        result_dataset.data = left_operand.data[list(result_dataset.components.keys())]
+        if left_operand.data is not None:
+            if is_from_component_assignment:
+                return DataComponent(
+                    name=right_operand,
+                    data_type=left_operand.components[right_operand].data_type,
+                    role=Role.MEASURE,
+                    nullable=left_operand.components[right_operand].nullable,
+                    data=left_operand.data[right_operand],
+                )
+            result_dataset.data = left_operand.data[list(result_dataset.components.keys())]
         return result_dataset
 
 
@@ -60,19 +69,19 @@ class Alias(Binary):
     It inherits from Binary class, and has the following class methods:
 
     Class methods:
-        Validate: Ensures the name given in the right operand is different from the name of the Dataset.
-        Evaluate: Checks if the data between both operators are the same.
+        Validate: Ensures the name given in the right operand is different from the
+        name of the Dataset. Evaluate: Checks if the data between both operators are the same.
     """
 
     @classmethod
-    def validate(cls, left_operand: Dataset, right_operand: str):
+    def validate(cls, left_operand: Dataset, right_operand: Union[str, Dataset]) -> Dataset:
         new_name = right_operand if isinstance(right_operand, str) else right_operand.name
         if new_name != left_operand.name and new_name in left_operand.get_components_names():
             raise SemanticError("1-3-1", alias=new_name)
         return Dataset(name=new_name, components=left_operand.components, data=None)
 
     @classmethod
-    def evaluate(cls, left_operand: Dataset, right_operand: str) -> Dataset:
+    def evaluate(cls, left_operand: Dataset, right_operand: Union[str, Dataset]) -> Dataset:
         result = cls.validate(left_operand, right_operand)
         result.data = left_operand.data
         return result
@@ -83,14 +92,16 @@ class Eval(Unary):
     It inherits from Unary class and has the following class methods
 
     Class methods:
-        Validate: checks if the external routine name is the same as the operand name, which must be a Dataset.
+        Validate: checks if the external routine name is the same as the operand name,
+        which must be a Dataset.
         Evaluate: Checks if the operand and the output is actually a Dataset.
 
     """
 
     @staticmethod
-    def _execute_query(query: str, dataset_names: List[str],
-                       data: Dict[str, pd.DataFrame]) -> pd.DataFrame:
+    def _execute_query(
+        query: str, dataset_names: List[str], data: Dict[str, pd.DataFrame]
+    ) -> pd.DataFrame:
         try:
             conn = sqlite3.connect(":memory:")
             try:
@@ -111,25 +122,33 @@ class Eval(Unary):
         return df_result
 
     @classmethod
-    def validate(cls, operands: Dict[str, Dataset], external_routine: ExternalRoutine,
-                 output: Dataset) -> Dataset:
+    def validate(  # type: ignore[override]
+        cls,
+        operands: Dict[str, Dataset],
+        external_routine: ExternalRoutine,
+        output: Dataset,
+    ) -> Dataset:
 
         empty_data_dict = {}
         for ds_name in external_routine.dataset_names:
             if ds_name not in operands:
-                raise ValueError(f"External Routine dataset {ds_name} "
-                                 f"is not present in Eval operands")
+                raise ValueError(
+                    f"External Routine dataset {ds_name} " f"is not present in Eval operands"
+                )
             empty_data = pd.DataFrame(
-                columns=[comp.name for comp in operands[ds_name].components.values()])
+                columns=[comp.name for comp in operands[ds_name].components.values()]
+            )
             empty_data_dict[ds_name] = empty_data
 
-        df = cls._execute_query(external_routine.query, external_routine.dataset_names,
-                                empty_data_dict)
+        df = cls._execute_query(
+            external_routine.query, external_routine.dataset_names, empty_data_dict
+        )
         component_names = [name for name in df.columns]
         for comp_name in component_names:
             if comp_name not in output.components:
-                raise SemanticError("1-1-1-10", op=cls.op, comp_name=comp_name,
-                                    dataset_name=df.name)
+                raise SemanticError(
+                    "1-1-1-10", op=cls.op, comp_name=comp_name, dataset_name=df.name
+                )
 
         for comp_name in output.components:
             if comp_name not in component_names:
@@ -140,15 +159,17 @@ class Eval(Unary):
         return output
 
     @classmethod
-    def evaluate(cls, operands: Dict[str, Dataset], external_routine: ExternalRoutine,
-                 output: Dataset) -> Dataset:
-        result = cls.validate(operands, external_routine, output)
-
-        operands_data_dict = {ds_name: operands[ds_name].data
-                              for ds_name in operands}
-
-        result.data = cls._execute_query(external_routine.query,
-                                         external_routine.dataset_names,
-                                         operands_data_dict)
-
+    def evaluate(  # type: ignore[override]
+        cls,
+        operands: Dict[str, Dataset],
+        external_routine: ExternalRoutine,
+        output: Dataset,
+    ) -> Dataset:
+        result: Dataset = cls.validate(operands, external_routine, output)
+        operands_data_dict = {ds_name: operands[ds_name].data for ds_name in operands}
+        result.data = cls._execute_query(
+            external_routine.query,
+            external_routine.dataset_names,
+            operands_data_dict,  # type: ignore[arg-type]
+        )
         return result