Refactor

luxonis · Jan 5, 2025 · c8d729c · c8d729c
1 parent 735e378
commit c8d729c
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 17 deletions.
diff --git a/bindings/python/src/pipeline/node/NodeBindings.cpp b/bindings/python/src/pipeline/node/NodeBindings.cpp
@@ -315,6 +315,14 @@ void NodeBindings::bind(pybind11::module& m, void* pCallstack) {
              DOC(dai, Node, Input, getParent))
         .def("getPossibleDatatypes", &Node::Input::getPossibleDatatypes, DOC(dai, Node, Input, getPossibleDatatypes))
         .def("setPossibleDatatypes", &Node::Input::setPossibleDatatypes, py::arg("types"), DOC(dai, Node, Input, setPossibleDatatypes))
+        .def("setPossibleDatatypes", [](Node::Input& input, const std::vector<std::tuple<DatatypeEnum, bool>>& types) {
+            std::vector<Node::DatatypeHierarchy> converted;
+            converted.reserve(types.size());
+            for(const auto& t : types) {
+                converted.emplace_back(std::get<0>(t), std::get<1>(t));
+            }
+            input.setPossibleDatatypes(converted);
+        }, py::arg("types"), DOC(dai, Node, Input, setPossibleDatatypes))
         .def("setWaitForMessage", &Node::Input::setWaitForMessage, py::arg("waitForMessage"), DOC(dai, Node, Input, setWaitForMessage))
         .def("getWaitForMessage", &Node::Input::getWaitForMessage, DOC(dai, Node, Input, getWaitForMessage))
         .def("setReusePreviousMessage", &Node::Input::setReusePreviousMessage, py::arg("reusePreviousMessage"), DOC(dai, Node, Input, setReusePreviousMessage))
@@ -339,7 +347,14 @@ void NodeBindings::bind(pybind11::module& m, void* pCallstack) {
              py::keep_alive<1, 0>())
         .def("getPossibleDatatypes", &Node::Output::getPossibleDatatypes, DOC(dai, Node, Output, getPossibleDatatypes))
         .def("setPossibleDatatypes", &Node::Output::setPossibleDatatypes, py::arg("types"), DOC(dai, Node, Output, setPossibleDatatypes))
-        .def("addPossibleDatatype", &Node::Output::addPossibleDatatype, py::arg("datatype"), py::arg("descendants") = true, DOC(dai, Node, Output, addPossibleDatatype))
+        .def("setPossibleDatatypes", [](Node::Output& output, const std::vector<std::tuple<DatatypeEnum, bool>>& types) {
+            std::vector<Node::DatatypeHierarchy> converted;
+            converted.reserve(types.size());
+            for(const auto& t : types) {
+                converted.emplace_back(std::get<0>(t), std::get<1>(t));
+            }
+            output.setPossibleDatatypes(converted);
+        }, py::arg("types"), DOC(dai, Node, Output, setPossibleDatatypes))
         .def("getParent",
              static_cast<const Node& (Node::Output::*)() const>(&Node::Output::getParent),
              py::return_value_policy::reference_internal,

diff --git a/include/depthai/pipeline/Node.hpp b/include/depthai/pipeline/Node.hpp
@@ -180,17 +180,12 @@ class Node : public std::enable_shared_from_this<Node> {
         /**
          * Get possible datatypes that can be sent
          */
-        std::vector<DatatypeHierarchy> getPossibleDatatypes() const {
-            return desc.types;
-        }
-
-        void setPossibleDatatypes(std::vector<DatatypeHierarchy> types) {
-            desc.types = std::move(types);
-        }
+        std::vector<DatatypeHierarchy> getPossibleDatatypes() const;
 
-        void addPossibleDatatype(DatatypeEnum datatype, bool descendants = true) {
-            desc.types.push_back(DatatypeHierarchy{datatype, descendants});
-        }
+        /**
+         * Set possible datatypes that can be sent
+         */
+        void setPossibleDatatypes(std::vector<DatatypeHierarchy> types);
 
         /**
          * Check if this output and given input are on the same pipeline.
@@ -386,13 +381,15 @@ class Node : public std::enable_shared_from_this<Node> {
          */
         bool getWaitForMessage() const;
 
-        std::vector<DatatypeHierarchy> getPossibleDatatypes() const {
-            return possibleDatatypes;
-        }
+        /**
+         * Get possible datatypes that can be received
+         */
+        std::vector<DatatypeHierarchy> getPossibleDatatypes() const;
 
-        void setPossibleDatatypes(std::vector<DatatypeHierarchy> types) {
-            possibleDatatypes = std::move(types);
-        }
+        /**
+         * Set possible datatypes that can be received
+         */
+        void setPossibleDatatypes(std::vector<DatatypeHierarchy> types);
 
         /**
          * Equivalent to setWaitForMessage but with inverted logic.

diff --git a/src/pipeline/Node.cpp b/src/pipeline/Node.cpp
@@ -731,4 +731,20 @@ std::vector<std::pair<Node::Input&, std::shared_ptr<Capability>>> Node::getRequi
     DAI_CHECK_V(false, "Node '{}' doesn't support node to node linking. Please link outputs <--> inputs manually.", getName());
 }
 
+void Node::Output::setPossibleDatatypes(std::vector<Node::DatatypeHierarchy> types) {
+    desc.types = std::move(types);
+}
+
+std::vector<Node::DatatypeHierarchy> Node::Output::getPossibleDatatypes() const {
+    return desc.types;
+}
+
+void Node::Input::setPossibleDatatypes(std::vector<Node::DatatypeHierarchy> types) {
+    possibleDatatypes = std::move(types);
+}
+
+std::vector<Node::DatatypeHierarchy> Node::Input::getPossibleDatatypes() const {
+    return possibleDatatypes;
+}
+
 }  // namespace dai