fix: sub module conflict error

Thriftpy · Nov 4, 2024 · 43126be · 43126be
1 parent e1d5a54
commit 43126be
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/thriftpy2/parser/__init__.py b/thriftpy2/parser/__init__.py
@@ -47,7 +47,7 @@ def load(path,
             sub_modules = thrift.__thrift_meta__["sub_modules"][:]
             for module in sub_modules:
                 if module not in sys.modules:
-                    sys.modules[module] = module
+                    sys.modules[module.__name__] = include_thrift
             if include_thrift.__name__ not in sys.modules:
                 include_thrifts.extend(include_thrift.__thrift_meta__["includes"])
     return thrift

diff --git a/thriftpy2/parser/parser.py b/thriftpy2/parser/parser.py
@@ -63,7 +63,7 @@ def p_include(p):
         path = os.path.join(include_dir, p[2])
         if os.path.exists(path):
             child_path = os.path.normpath(
-                os.path.dirname(str(thrift.__name__).rstrip("_thrift").replace(".", "/")) + "/" + p[2])
+                os.path.dirname(remove_suffix(str(thrift.__name__), "_thrift").replace(".", "/")) + "/" + p[2])
 
             child_path = child_path.lstrip("/")
 
@@ -72,8 +72,8 @@ def p_include(p):
                                     ".").replace(
                 ".thrift", "_thrift")
 
-            child = parse(path)
-            setattr(thrift, child.__name__, child)
+            child = parse(path, module_name=child_module_name)
+            setattr(thrift, remove_suffix(child.__name__, "_thrift"), child)
             _add_thrift_meta('includes', child)
             _add_thrift_meta('sub_modules', types.ModuleType(child_module_name))
             return
@@ -960,3 +960,8 @@ def _get_ttype(inst, default_ttype=None):
     if hasattr(inst, '__dict__') and '_ttype' in inst.__dict__:
         return inst.__dict__['_ttype']
     return default_ttype
+
+def remove_suffix(s, suffix):
+    if s.endswith(suffix):
+        return s[:-len(suffix)]
+    return s