Fix jinja

Quantco · Nov 22, 2024 · 3f25d7e · 3f25d7e
1 parent c9de7ec
commit 3f25d7e
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/tools/templates/class.jinja2 b/tools/templates/class.jinja2
@@ -14,11 +14,11 @@ class _{{ schema.name }}(StandardNode):
     {% for input in schema.inputs %}
         {{ input.name }}: {%
             if is_optional(input)
-                %}Optional[VarInfo]{%
+                %}Optional[_VarInfo]{%
             elif is_variadic(input)
-                %}Sequence[VarInfo]{%
+                %}Sequence[_VarInfo]{%
             else
-                %}VarInfo{%
+                %}_VarInfo{%
             endif %}
 
     {% endfor %}
@@ -33,11 +33,11 @@ class _{{ schema.name }}(StandardNode):
     {% for output in schema.outputs %}
         {{ output.name }}: {%
             if is_optional(output)
-                %}Optional[VarInfo]{%
+                %}Optional[_VarInfo]{%
             elif is_variadic(output)
-                %}Sequence[VarInfo]{%
+                %}Sequence[_VarInfo]{%
             else
-                %}VarInfo{%
+                %}_VarInfo{%
             endif %}
 
     {% endfor %}

diff --git a/tools/templates/preamble.jinja2 b/tools/templates/preamble.jinja2
@@ -14,7 +14,7 @@ from typing import cast as typing_cast
 import numpy as np
 import numpy.typing as npt
 
-from spox._var import Var, VarInfo, result_type, unwrap_vars, get_value
+from spox._var import Var, _VarInfo, result_type, unwrap_vars, get_value
 from spox._fields import BaseAttributes, BaseInputs, BaseOutputs
 from spox._attributes import (
     AttrDtype,