diff --git a/src/bindings/python/src/openvino/_ov_api.py b/src/bindings/python/src/openvino/_ov_api.py index 9b8f5f3ca06d71..5c08537e8bb1a3 100644 --- a/src/bindings/python/src/openvino/_ov_api.py +++ b/src/bindings/python/src/openvino/_ov_api.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 from types import TracebackType -from typing import Any, Iterable, Union, Optional, Dict, Type +from typing import Any, Iterable, Union, Optional, Dict, Type, List from pathlib import Path @@ -12,8 +12,7 @@ from openvino._pyopenvino import CompiledModel as CompiledModelBase from openvino._pyopenvino import AsyncInferQueue as AsyncInferQueueBase from openvino._pyopenvino import Op as OpBase -from openvino._pyopenvino import Tensor -from openvino._pyopenvino import Node +from openvino._pyopenvino import Node, Output, Tensor from openvino.utils.data_helpers import ( OVDict, @@ -24,7 +23,7 @@ class Op(OpBase): - def __init__(self, py_obj, inputs=None) -> None: + def __init__(self, py_obj: "Op", inputs: List[Union[Node, Output]] = None) -> None: super().__init__(py_obj) self._update_type_info() if inputs is not None: diff --git a/src/bindings/python/tests/test_graph/test_custom_op.py b/src/bindings/python/tests/test_graph/test_custom_op.py index a2e5588d97836c..63ac8e8f4b1d78 100644 --- a/src/bindings/python/tests/test_graph/test_custom_op.py +++ b/src/bindings/python/tests/test_graph/test_custom_op.py @@ -87,7 +87,7 @@ class CustomOpWithAttribute(Op): class_type_info = DiscreteTypeInfo("CustomOpWithAttribute", "extension") def __init__(self, inputs=None, attrs=None): - super().__init__(self, inputs) + super().__init__(self) if attrs is not None or inputs is not None: self._attrs = attrs self.set_arguments(inputs) @@ -239,9 +239,9 @@ def test_op_extension(prepared_paths): param1 = ops.parameter(Shape(input_shape), dtype=np.float32, name="data1") param2 = ops.parameter(Shape(input_shape), dtype=np.float32, name="data2") - custom_simple = CustomSimpleOp(inputs=[param1, param2]) + custom_simple = CustomSimpleOp([param1, param2]) custom_simple.set_friendly_name("test_add") - custom_with_attribute = CustomSimpleOpWithAttribute(inputs=[custom_simple], value_str="test_attribute") + custom_with_attribute = CustomSimpleOpWithAttribute([custom_simple], value_str="test_attribute") custom_add = CustomAdd(inputs=[custom_with_attribute]) res = ops.result(custom_add, name="result") simple_model = Model(res, [param1, param2], "SimpleModel")