resolve comments

Author: Priyadarshini Piramanayagam

nisystemlink/clients/testmonitor/utilities/_dataframe_utilities.py

@@ -1,22 +1,25 @@
 from typing import Any, Callable, Dict, List, Optional
 
 import pandas as pd
-from nisystemlink.clients.testmonitor.models import Result, Step, StepProjection
+from nisystemlink.clients.testmonitor.models import (
+    Measurement,
+    Result,
+    Step,
+    StepProjection,
+)
 from nisystemlink.clients.testmonitor.utilities.constants import DataFrameHeaders
 
 
-def is_step_data_with_name_and_measurement(data: Dict[str, Any]) -> bool:
-    """Checks if a step data parameter is a measurement data by ensuring it 
-    has both 'name' and 'measurement' fields.
+def is_step_data_with_name_and_measurement(measurement: Measurement) -> bool:
+    """Checks if a measurement data has both 'name' and 'measurement' fields.
 
     Args:
-        measurement: A dictionary containing measurement data.
+        measurement: A measurement data object
 
     Returns:
         bool: True if the measurement has both 'name' and 'measurement' fields, False otherwise.
     """
-    required_measurement_fields = ["name", "measurement"]
-    return set(required_measurement_fields).issubset(set(data.keys()))
+    return measurement.name is not None and measurement.measurement is not None
 
 
 def convert_results_to_dataframe(
@@ -55,7 +58,7 @@ def convert_results_to_dataframe(
 def convert_steps_to_dataframe(
     steps: List[Step],
     is_measurement_data_parameter: Optional[
-        Callable[[Dict[str, Any]], bool]
+        Callable[[Measurement], bool]
     ] = is_step_data_with_name_and_measurement,
 ) -> pd.DataFrame:
     """Converts a list of steps into a normalized dataframe.
@@ -64,8 +67,8 @@ def convert_steps_to_dataframe(
         steps: A list of steps.
         is_measurement_data_parameter: Optional callback function that checks if a step data parameter is a
             measurement so that only those are included in the returned dataframe. The method takes
-            a dictionary as input and returns a boolean value.
-            The default behavior is to consider only parameters that have both 'name' and 'measurement'
+            a measurement as input and returns a boolean value.
+            The default behavior is to consider only measurement data that have both 'name' and 'measurement'
             fields as measurement parameters.
             If none of the step data parameters have the required fields, the data.parameters will not
             appear in the dataframe.
@@ -184,59 +187,53 @@ def __is_property_header(header: str) -> bool:
 
 def __convert_steps_to_dict(
     steps: List[Step],
-    is_measurement_data_parameter: Optional[Callable[[Dict[str, Any]], bool]],
+    is_measurement_data_parameter: Optional[Callable[[Measurement], bool]],
 ) -> List[Dict[str, Any]]:
     """Converts a list of steps to dictionaries, excluding None values.
 
     Args:
         steps: A list of steps.
-        is_measurement_data_parameter: Optional callback function to check if a step
-            data parameter has the required fields.
+        is_measurement_data_parameter: Optional callback function that checks if a step data
+            parameter is a measurement so that only those are included in the returned dataframe.
 
     Returns:
         List[Dict[str, Any]]: A list of dictionaries containing step information.
     """
     steps_dict = []
     for step in steps:
-        single_step_dict = step.dict(exclude_none=True)
+        __filter_data_parameters(step, is_measurement_data_parameter)
 
-        __filter_data_parameters(single_step_dict, step, is_measurement_data_parameter)
+        single_step_dict = step.dict(exclude_none=True)
         __normalize_inputs_outputs(single_step_dict, step)
         __normalize_step_status(single_step_dict)
         steps_dict.append(single_step_dict)
     return steps_dict
 
 
 def __filter_data_parameters(
-    step_dict: Dict[str, Any],
     step: Step,
-    is_measurement_data_parameter: Optional[Callable[[Dict[str, Any]], bool]],
+    is_measurement_data_parameter: Optional[Callable[[Measurement], bool]],
 ) -> None:
     """Gets data parameters from the step dictionary and filters it based on the callback function.
 
     Args:
-        step_dict: A dictionary with step information.
         step: A Step object containing data parameters.
         is_measurement_data_parameter: Optional callback function to check if a measurement is valid. The method takes
             a dictionary as input and returns a boolean value. The default behavior is to consider only parameters
             that have both 'name' and 'measurement' fields as measurement parameters.
 
     Returns:
-        None: The function modifies step_dict in place with the filtered step data parameters.
+        None: The function modifies step parameters in place with filtered data parameters.
     """
-    valid_measurement_parameters = []
-    if (
-        is_measurement_data_parameter is not None
-        and step.data is not None
-        and step.data.parameters is not None
-    ):
-        for parameter in step_dict["data"]["parameters"]:
+    if step.data and step.data.parameters and is_measurement_data_parameter is not None:
+        valid_measurement_parameters = []
+        for measurement in step.data.parameters:
             if is_measurement_data_parameter and is_measurement_data_parameter(
-                parameter
+                measurement
             ):
-                valid_measurement_parameters.append(parameter)
+                valid_measurement_parameters.append(measurement)
 
-        step_dict["data"]["parameters"] = valid_measurement_parameters
+        step.data.parameters = valid_measurement_parameters
 
 
 def __normalize_inputs_outputs(
@@ -330,9 +327,7 @@ def __group_step_columns(dataframe_columns: List[str]) -> List[str]:
         StepProjection.DATA.lower(),
         StepProjection.PROPERTIES.lower(),
     ]
-    grouped_columns: Dict[str, List[str]] = {
-        category: [] for category in CATEGORY_KEYS
-    }
+    grouped_columns: Dict[str, List[str]] = {category: [] for category in CATEGORY_KEYS}
     for column in dataframe_columns:
         column_lower = column.lower()
         key = next(