From a30a03e80bc2512e796ff1002e18553a1922d8dc Mon Sep 17 00:00:00 2001
From: ethframe <ethframe@gmail.com>
Date: Wed, 14 Jul 2021 21:17:50 +0300
Subject: [PATCH] Add package prefix to imports in generated code (#313)

* Add package prefix to imports in generated code

Packages generated with thrift refers to each other as if they are top
level packages. This prevents imports from working when generated
packages are placed inside some other package. Currently jaeger_client
relies on hacking sys.path, but this doesn't work properly if the
library is packed inside an application built using PyInstaller.
Another way to get imports to work is to rewrite them after
generation. This can be done easily since the build process already
involves an awk script that performs a couple of replacements.

Signed-off-by: Andrey Nikitin <ethframe@gmail.com>

* Remove unused import

Signed-off-by: Andrey Nikitin <ethframe@gmail.com>
---
 Makefile                                 | 3 ++-
 jaeger_client/__init__.py                | 7 -------
 jaeger_client/thrift_gen/agent/Agent.py  | 8 ++++----
 jaeger_client/thrift_gen/agent/ttypes.py | 4 ++--
 thrift-gen-fix.awk                       | 3 +++
 5 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/Makefile b/Makefile
index 9609d1a8..74b75c79 100644
--- a/Makefile
+++ b/Makefile
@@ -76,6 +76,7 @@ shell:
 
 # Generate jaeger thrifts
 THRIFT_GEN_DIR=jaeger_client/thrift_gen
+THRIFT_PACKAGE_PREFIX=jaeger_client.thrift_gen
 THRIFT_VER=0.9.3
 THRIFT_IMG=thrift:$(THRIFT_VER)
 THRIFT_PY_ARGS=new_style,tornado
@@ -100,7 +101,7 @@ thrift: idl-submodule thrift-image
 	set -e; \
 	for f in $$(find ${THRIFT_GEN_DIR} -iname '*.py'); do \
 	  echo fixing $$f; \
-	  awk -f thrift-gen-fix.awk $$f > tmp; \
+	  awk -f thrift-gen-fix.awk package_prefix=${THRIFT_PACKAGE_PREFIX} $$f > tmp; \
 	  mv tmp $$f; \
 	done
 
diff --git a/jaeger_client/__init__.py b/jaeger_client/__init__.py
index c523ab79..58462537 100644
--- a/jaeger_client/__init__.py
+++ b/jaeger_client/__init__.py
@@ -14,13 +14,6 @@
 
 from __future__ import absolute_import
 
-import sys
-
-# This is because thrift for python doesn't have 'package_prefix'.
-# The thrift compiled libraries refer to each other relative to their subdir.
-import jaeger_client.thrift_gen as modpath
-sys.path.append(modpath.__path__[0])
-
 __version__ = '4.5.1.dev0'
 
 from .tracer import Tracer  # noqa
diff --git a/jaeger_client/thrift_gen/agent/Agent.py b/jaeger_client/thrift_gen/agent/Agent.py
index c7e28f93..a495c771 100644
--- a/jaeger_client/thrift_gen/agent/Agent.py
+++ b/jaeger_client/thrift_gen/agent/Agent.py
@@ -153,7 +153,7 @@ class emitZipkinBatch_args(object):
 
   thrift_spec = (
     None, # 0
-    (1, TType.LIST, 'spans', (TType.STRUCT,(zipkincore.ttypes.Span, zipkincore.ttypes.Span.thrift_spec)), None, ), # 1
+    (1, TType.LIST, 'spans', (TType.STRUCT,(jaeger_client.thrift_gen.zipkincore.ttypes.Span, jaeger_client.thrift_gen.zipkincore.ttypes.Span.thrift_spec)), None, ), # 1
   )
 
   def __init__(self, spans=None,):
@@ -173,7 +173,7 @@ def read(self, iprot):
           self.spans = []
           (_etype3, _size0) = iprot.readListBegin()
           for _i4 in xrange(_size0):
-            _elem5 = zipkincore.ttypes.Span()
+            _elem5 = jaeger_client.thrift_gen.zipkincore.ttypes.Span()
             _elem5.read(iprot)
             self.spans.append(_elem5)
           iprot.readListEnd()
@@ -227,7 +227,7 @@ class emitBatch_args(object):
 
   thrift_spec = (
     None, # 0
-    (1, TType.STRUCT, 'batch', (jaeger.ttypes.Batch, jaeger.ttypes.Batch.thrift_spec), None, ), # 1
+    (1, TType.STRUCT, 'batch', (jaeger_client.thrift_gen.jaeger.ttypes.Batch, jaeger_client.thrift_gen.jaeger.ttypes.Batch.thrift_spec), None, ), # 1
   )
 
   def __init__(self, batch=None,):
@@ -244,7 +244,7 @@ def read(self, iprot):
         break
       if fid == 1:
         if ftype == TType.STRUCT:
-          self.batch = jaeger.ttypes.Batch()
+          self.batch = jaeger_client.thrift_gen.jaeger.ttypes.Batch()
           self.batch.read(iprot)
         else:
           iprot.skip(ftype)
diff --git a/jaeger_client/thrift_gen/agent/ttypes.py b/jaeger_client/thrift_gen/agent/ttypes.py
index 9529cf29..bd68a628 100644
--- a/jaeger_client/thrift_gen/agent/ttypes.py
+++ b/jaeger_client/thrift_gen/agent/ttypes.py
@@ -9,8 +9,8 @@
 from six.moves import xrange
 
 from thrift.Thrift import TType, TMessageType, TException, TApplicationException
-import jaeger.ttypes
-import zipkincore.ttypes
+import jaeger_client.thrift_gen.jaeger.ttypes
+import jaeger_client.thrift_gen.zipkincore.ttypes
 
 
 from thrift.transport import TTransport
diff --git a/thrift-gen-fix.awk b/thrift-gen-fix.awk
index 3286a0bd..b9952853 100644
--- a/thrift-gen-fix.awk
+++ b/thrift-gen-fix.awk
@@ -11,5 +11,8 @@ BEGIN {six=0}
 {
     gsub(/from ttype/, "from .ttype", $0);
     gsub(/self.__dict__.iteritems\(\)/, "six.iteritems(self.__dict__)", $0);
+    if (package_prefix) {
+        $0 = gensub(/[[:alnum:]_]+\.ttypes/, package_prefix ".\\0", "g", $0);
+    }
     print
 }
\ No newline at end of file