add types

openvinotoolkit · Jan 10, 2025 · 948d847 · 948d847
1 parent 5a5081e
commit 948d847
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git 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
@@ -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")