Add domain as an optional parameter for make_node function (#1588)

yskim1501 · bddppq · commit f86e951d6957 · 2018-11-07T11:49:08.000-07:00
* add domain as an optional parameter for creating NodeProto

* at least two space for comments

* add test case

* add suggested comment

* modify comment

* explicitly check for None
diff --git a/onnx/helper.py b/onnx/helper.py
@@ -23,6 +23,7 @@ def make_node(
         outputs,  # type: Sequence[Text]
         name=None,  # type: Optional[Text]
         doc_string=None,  # type: Optional[Text]
+        domain=None,  # type: Optional[Text]
         **kwargs  # type: Any
 ):  # type: (...) -> NodeProto
     """Construct a NodeProto.
@@ -33,6 +34,8 @@ def make_node(
         outputs (list of string): list of output names
         name (string, default None): optional unique identifier for NodeProto
         doc_string (string, default None): optional documentation string for NodeProto
+        domain (string, default None): optional domain for NodeProto.
+            If it's None, we will just use default domain (which is empty)
         **kwargs (dict): the attributes of the node.  The acceptable values
             are documented in :func:`make_attribute`.
     """
@@ -45,6 +48,8 @@ def make_node(
         node.name = name
     if doc_string:
         node.doc_string = doc_string
+    if domain is not None:
+        node.domain = domain
     if kwargs:
         node.attribute.extend(
             make_attribute(key, value)
diff --git a/onnx/test/helper_test.py b/onnx/test/helper_test.py
@@ -202,6 +202,11 @@ def test_node_with_arg(self):  # type: () -> None
             node_def.attribute[0],
             helper.make_attribute("arg_value", 1))
 
+    def test_node_domain(self):  # type: () -> None
+        node_def = helper.make_node(
+            "Relu", ["X"], ["Y"], name="test", doc_string="doc", domain="test.domain")
+        self.assertEqual(node_def.domain, "test.domain")
+
     def test_graph(self):  # type: () -> None
         node_def1 = helper.make_node(
             "Relu", ["X"], ["Y"])
