From b1e80647b1391333fcf2e4bddc736171330faeb2 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 8c91c010f..9961242c8 100644
--- a/shared/src/main/scala/io/kaitai/struct/translators/PythonTranslator.scala
+++ b/shared/src/main/scala/io/kaitai/struct/translators/PythonTranslator.scala
@@ -63,8 +63,11 @@ class PythonTranslator(provider: TypeProvider, importList: ImportList, config: R
     }
     s"${PythonCompiler.types2class(enumSpec.name, isExternal)}.$label"
   }
-  override def doEnumById(enumSpec: EnumSpec, id: String): String =
+  override def doEnumById(enumSpec: EnumSpec, id: String): String = {
+    importList.add(s"from kaitaistruct import ${PythonCompiler.kstreamName}")
+
     s"${PythonCompiler.kstreamName}.resolve_enum(${PythonCompiler.types2class(enumSpec.name, enumSpec.isExternal(provider.nowClass))}, $id)"
+  }
 
   override def booleanOp(op: Ast.boolop) = op match {
     case Ast.boolop.Or => "or"
@@ -103,16 +106,25 @@ class PythonTranslator(provider: TypeProvider, importList: ImportList, config: R
 
   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 =