From 10ce7b6b6c55feadc408932f53eddf5c7491c8fb Mon Sep 17 00:00:00 2001
From: Mingun <alexander_sergey@mail.ru>
Date: Mon, 8 Apr 2024 23:10:53 +0500
Subject: [PATCH] Python, Construct: add missing imports of KaitaiStream

Fixes tests:
- Python:
  - ExprBytesNonLiteral
- Construct:
  - EnumOfValueInst
  - ExprBytesNonLiteral
  - ExprBytesOps
  - ExprEnum
---
 .../struct/translators/PythonTranslator.scala | 20 +++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/shared/src/main/scala/io/kaitai/struct/translators/PythonTranslator.scala b/shared/src/main/scala/io/kaitai/struct/translators/PythonTranslator.scala
index 2764e0bb0..0f6daa988 100644
--- a/shared/src/main/scala/io/kaitai/struct/translators/PythonTranslator.scala
+++ b/shared/src/main/scala/io/kaitai/struct/translators/PythonTranslator.scala
@@ -58,8 +58,11 @@ class PythonTranslator(provider: TypeProvider, importList: ImportList) extends B
 
   override def doEnumByLabel(enumTypeAbs: List[String], label: String): String =
     s"${PythonCompiler.types2class(enumTypeAbs)}.$label"
-  override def doEnumById(enumTypeAbs: List[String], id: String): String =
+  override def doEnumById(enumTypeAbs: List[String], id: String): String = {
+    importList.add(s"from kaitaistruct import ${PythonCompiler.kstreamName}")
+
     s"${PythonCompiler.kstreamName}.resolve_enum(${PythonCompiler.types2class(enumTypeAbs)}, $id)"
+  }
 
   override def booleanOp(op: Ast.boolop) = op match {
     case Ast.boolop.Or => "or"
@@ -98,16 +101,25 @@ class PythonTranslator(provider: TypeProvider, importList: ImportList) extends B
 
   override def bytesLength(value: Ast.expr): String =
     s"len(${translate(value)})"
-  override def bytesSubscript(container: Ast.expr, idx: Ast.expr): String =
+  override def bytesSubscript(container: Ast.expr, idx: Ast.expr): String = {
+    importList.add(s"from kaitaistruct import ${PythonCompiler.kstreamName}")
+
     s"${PythonCompiler.kstreamName}.byte_array_index(${translate(container)}, ${translate(idx)})"
+  }
   override def bytesFirst(a: Ast.expr): String =
     bytesSubscript(a, Ast.expr.IntNum(0))
   override def bytesLast(a: Ast.expr): String =
     bytesSubscript(a, Ast.expr.IntNum(-1))
-  override def bytesMin(b: Ast.expr): String =
+  override def bytesMin(b: Ast.expr): String = {
+    importList.add(s"from kaitaistruct import ${PythonCompiler.kstreamName}")
+
     s"${PythonCompiler.kstreamName}.byte_array_min(${translate(b)})"
-  override def bytesMax(b: Ast.expr): String =
+  }
+  override def bytesMax(b: Ast.expr): String = {
+    importList.add(s"from kaitaistruct import ${PythonCompiler.kstreamName}")
+
     s"${PythonCompiler.kstreamName}.byte_array_max(${translate(b)})"
+  }
 
 
   override def strLength(value: Ast.expr): String =