diff --git a/backend_py/primary/primary/services/flow_network_assembler/_group_tree_dataframe_model.py b/backend_py/primary/primary/services/flow_network_assembler/_group_tree_dataframe_model.py index c0d3e797f..969579f84 100644 --- a/backend_py/primary/primary/services/flow_network_assembler/_group_tree_dataframe_model.py +++ b/backend_py/primary/primary/services/flow_network_assembler/_group_tree_dataframe_model.py @@ -171,7 +171,7 @@ def _create_branch_node_list(self, terminal_node: str) -> List[str]: current_parents = [terminal_node] while len(current_parents) > 0: # Find all indexes matching the current parents - children_indices = set([i for i, x in enumerate(parents_array) if x in current_parents]) + children_indices = {i for i, x in enumerate(parents_array) if x in current_parents} # Find all children of the current parents children = nodes_array[list(children_indices)] diff --git a/backend_py/primary/primary/services/flow_network_assembler/_utils.py b/backend_py/primary/primary/services/flow_network_assembler/_utils.py index 27cee849f..627f67ed0 100644 --- a/backend_py/primary/primary/services/flow_network_assembler/_utils.py +++ b/backend_py/primary/primary/services/flow_network_assembler/_utils.py @@ -157,7 +157,7 @@ def create_sumvec_from_datatype_nodename_and_keyword( try: if keyword == "WELSPECS": datatype_ecl = WELL_DATATYPE_VECTOR_MAP[datatype] - elif keyword in TreeType._member_names_: + elif keyword in [t.value for t in TreeType]: datatype_ecl = GROUPTYPE_DATATYPE_VECTORS_MAP[TreeType[keyword]][datatype] except KeyError as exc: error = ( diff --git a/backend_py/primary/primary/services/flow_network_assembler/flow_network_assembler.py b/backend_py/primary/primary/services/flow_network_assembler/flow_network_assembler.py index edbbd8487..14a5993a1 100644 --- a/backend_py/primary/primary/services/flow_network_assembler/flow_network_assembler.py +++ b/backend_py/primary/primary/services/flow_network_assembler/flow_network_assembler.py @@ -205,7 +205,6 @@ async def fetch_and_initialize_async(self) -> None: # Get all vectors of interest existing in the summary data vectors_of_interest = _utils.get_all_vectors_of_interest_for_tree(group_tree_df_model) vectors_of_interest = vectors_of_interest & self._all_vectors - # TODO: do we need to return to a list? # Has any water injection or gas injection vectors among vectors of interest has_wi_vectors = False @@ -886,8 +885,8 @@ def _get_data_for_summary_vector_info( """Extracts and formats whatever data is available for a given summary vector. If the name is not present in the summary, a list of NaN values is given instead""" if sumvec_name in summary_col_set: return smry_for_grouptree_sorted_by_date[sumvec_name].to_numpy().round(2) - else: - return list([np.nan] * number_of_dates) + + return list([np.nan] * number_of_dates) def _is_valid_node_type(node_classification: NodeClassification, valid_node_types: set[NodeType]) -> bool: diff --git a/backend_py/primary/primary/services/flow_network_assembler/flow_network_types.py b/backend_py/primary/primary/services/flow_network_assembler/flow_network_types.py index 06a478780..22c2d6d16 100644 --- a/backend_py/primary/primary/services/flow_network_assembler/flow_network_types.py +++ b/backend_py/primary/primary/services/flow_network_assembler/flow_network_types.py @@ -1,7 +1,5 @@ -from pydantic import BaseModel - - from typing import Dict, List, Literal +from pydantic import BaseModel class NetworkNode(BaseModel):