aimclub · maypink · Nov 18, 2022 · Nov 15, 2022 · Nov 15, 2022 · Nov 15, 2022
diff --git a/fedot/core/dag/graph.py b/fedot/core/dag/graph.py
@@ -103,6 +103,34 @@ def get_edges(self) -> Sequence[Tuple[GraphNode, GraphNode]]:
         """
         raise NotImplementedError()
 
+    def get_nodes_by_name(self, name: str) -> List[GraphNode]:
+        """Returns list of nodes with the required ``name``
+
+        Args:
+            name: name to filter by
+
+        Returns:
+            list: relevant nodes (empty if there are no such nodes)
+        """
+
+        appropriate_nodes = filter(lambda x: x.name == name, self.nodes)
+
+        return list(appropriate_nodes)
+
+    def get_node_by_uid(self, uid: str) -> Optional[GraphNode]:
+        """Returns node with the required ``uid``
+
+        Args:
+            uid: uid of node to filter by
+
+        Returns:
+            Optional[Node]: relevant node (None if there is no such node)
+        """
+
+        appropriate_nodes = list(filter(lambda x: x.uid == uid, self.nodes))
+
+        return appropriate_nodes[0] if appropriate_nodes else None
+
     @abstractmethod
     def __eq__(self, other_graph: 'Graph') -> bool:
         """Compares this graph with the ``other_graph``

diff --git a/fedot/core/dag/graph_node.py b/fedot/core/dag/graph_node.py
@@ -1,6 +1,7 @@
 from abc import ABC, abstractmethod
 from copy import copy
 from typing import List, Optional, Iterable
+from uuid import uuid4
 
 
 class GraphNode(ABC):
@@ -9,6 +10,8 @@ class GraphNode(ABC):
     Provides interface for getting and modifying the parent nodes
     and recursive description based on all preceding nodes.
     """
+    def __init__(self):
+        self.uid = str(uuid4())
 
     @property
     @abstractmethod
@@ -30,6 +33,12 @@ def nodes_from(self, nodes: Optional[Iterable['GraphNode']]):
         """
         pass
 
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """ Str name of this graph node """
+        pass
+
     @abstractmethod
     def __str__(self) -> str:
         """Returns short node type description

diff --git a/fedot/core/dag/linked_graph_node.py b/fedot/core/dag/linked_graph_node.py
@@ -28,7 +28,6 @@ def __init__(self, content: Union[dict, str],
 
         self.content = content
         self._nodes_from = UniqueList(nodes_from or ())
-        self.uid = str(uuid4())
 
         super().__init__()
 
@@ -40,11 +39,15 @@ def nodes_from(self) -> List['LinkedGraphNode']:
     def nodes_from(self, nodes: Optional[Iterable['LinkedGraphNode']]):
         self._nodes_from = UniqueList(nodes)
 
+    @property
+    def name(self) -> str:
+        return str(self.content.get('name'))
+
     def __hash__(self) -> int:
         return hash(self.uid)
 
     def __str__(self) -> str:
-        return str(self.content['name'])
+        return str(self.content.get('name'))
 
     def __repr__(self) -> str:
         return self.__str__()

diff --git a/fedot/core/pipelines/node.py b/fedot/core/pipelines/node.py
@@ -112,6 +112,11 @@ def update_params(self):
         updated_parameters = {**self.parameters, **changed_parameters}
         self.parameters = updated_parameters
 
+    @property
+    def name(self) -> str:
+        """ Returns str name of operation """
+        return self.operation.operation_type
+
     @property
     def operation(self) -> Operation:
         """Returns node operation object

diff --git a/fedot/core/pipelines/pipeline.py b/fedot/core/pipelines/pipeline.py
@@ -371,23 +371,6 @@ def show(self, save_path: Optional[Union[PathLike, str]] = None, engine: Optiona
         PipelineVisualizer(self).visualise(save_path, engine, node_color, dpi, node_size_scale, font_size_scale)
 
 
-def nodes_with_operation(pipeline: Pipeline, operation_name: str) -> List[Node]:
-    """Returns list of nodes with the required ``operation_name``
-
-    Args:
-        pipeline: pipeline to process
-        operation_name: name of the operation to filter by
-
-    Returns:
-        list: relevant nodes (empty if there are no such nodes)
-    """
-
-    # Check if model has decompose operations
-    appropriate_nodes = filter(lambda x: x.operation.operation_type == operation_name, pipeline.nodes)
-
-    return list(appropriate_nodes)
-
-
 def _graph_nodes_to_pipeline_nodes(operator: LinkedGraph, nodes: Sequence[Node]):
     """
     Method to update nodes type after performing some action on the pipeline

diff --git a/fedot/core/pipelines/pipeline_builder.py b/fedot/core/pipelines/pipeline_builder.py
@@ -53,6 +53,7 @@ def add_sequence(self, *operation_type: OperationType, branch_idx: int = 0):
 
         :param operation_type: operations for new nodes, either as an operation name
             or as a tuple of operation name and operation parameters.
+        :param branch_idx: index of the branch for branching its tip
         """
         for operation in operation_type:
             operation, params = self._unpack_operation(operation)

diff --git a/fedot/core/pipelines/verification_rules.py b/fedot/core/pipelines/verification_rules.py
@@ -2,7 +2,7 @@
 
 from fedot.core.operations.model import Model
 from fedot.core.pipelines.node import PrimaryNode
-from fedot.core.pipelines.pipeline import Pipeline, nodes_with_operation
+from fedot.core.pipelines.pipeline import Pipeline
 from fedot.core.repository.dataset_types import DataTypesEnum
 from fedot.core.repository.operation_types_repository import OperationTypesRepository, get_operations_for_task, \
     atomized_model_type
@@ -180,8 +180,7 @@ def has_no_conflicts_in_decompose(pipeline: Pipeline):
     """ The function checks whether the 'class_decompose' or 'decompose' operation has two ancestors """
 
     for decomposer in ['decompose', 'class_decompose']:
-        decompose_nodes = nodes_with_operation(pipeline,
-                                               decomposer)
+        decompose_nodes = pipeline.get_nodes_by_name(decomposer)
         if len(decompose_nodes) != 0:
             # Launch check decomposers
             __check_decomposer_has_two_parents(nodes_to_check=decompose_nodes)

diff --git a/test/unit/pipelines/test_decompose_pipelines.py b/test/unit/pipelines/test_decompose_pipelines.py
@@ -6,7 +6,7 @@
 from examples.advanced.decompose.refinement_forecast_example import get_refinement_pipeline
 from fedot.core.data.data import InputData
 from fedot.core.pipelines.node import PrimaryNode, SecondaryNode
-from fedot.core.pipelines.pipeline import Pipeline, nodes_with_operation
+from fedot.core.pipelines.pipeline import Pipeline
 from fedot.core.repository.dataset_types import DataTypesEnum
 from fedot.core.repository.tasks import Task, TaskTypesEnum
 from fedot.preprocessing.preprocessing import DataPreprocessor
@@ -159,13 +159,13 @@ def test_order_by_data_flow_len_correct():
         pipeline.fit(input_data)
 
         # Get one node with decompose operation in it
-        decompose_nodes = nodes_with_operation(pipeline, 'class_decompose')
+        decompose_nodes = pipeline.get_nodes_by_name('class_decompose')
         decompose_node = decompose_nodes[0]
         # Predict from decompose must be the same as predict from Data parent
         dec_output = decompose_node.predict(input_data)
 
         # Get data parent operation for node
-        data_node = nodes_with_operation(pipeline, data_operation)[0]
+        data_node = pipeline.get_nodes_by_name(data_operation)[0]
         data_output = data_node.predict(input_data)
 
         if tuple(data_output.predict.shape) != tuple(dec_output.predict.shape):

diff --git a/test/unit/pipelines/test_pipeline.py b/test/unit/pipelines/test_pipeline.py
@@ -436,3 +436,23 @@ def test_ts_forecasting_pipeline_with_poly_features():
     pipeline.fit(train_data)
     prediction = pipeline.predict(test_data)
     assert prediction is not None
+
+
+def test_get_nodes_with_operation():
+    pipeline = pipeline_first()
+    actual_nodes = pipeline.get_nodes_by_name(name='rf')
+    expected_nodes = [pipeline.nodes[2], pipeline.nodes[-1]]
+
+    assert (actual is expected for actual, expected in zip(actual_nodes, expected_nodes))
+
+
+def test_get_node_with_uid():
+    pipeline = pipeline_first()
+
+    uid_of_first_node = pipeline.nodes[0].uid
+    actual_node = pipeline.get_node_by_uid(uid=uid_of_first_node)
+    expected_node = pipeline.nodes[0]
+    assert actual_node is expected_node
+
+    uid_of_non_existent_node = '123456789'
+    assert pipeline.get_node_by_uid(uid=uid_of_non_existent_node) is None
diff --git a/test/unit/serialization/mocks/history_mocks.py b/test/unit/serialization/mocks/history_mocks.py
@@ -8,13 +8,18 @@ class CustomMockNode(GraphNode):
     def __init__(self, content: dict = None, nodes_from: list = None):
         self._content = content
         self._nodes_from = nodes_from or []
+        super().__init__()
 
     @property
     def nodes_from(self):
         return self._nodes_from
 
-    def __str__(self):
-        return self._content['name']
+    @property
+    def name(self) -> str:
+        return self._content.get('name')
+
+    def __str__(self) -> str:
+        return self._content.get('name')
 
 
 class CustomMockGraph(Graph):