Add submodule to sys.path when loading child idl file. (#249)

* add submodule to sys.path when loading module * add sub modules to sys.modules recursively
Thriftpy · Mar 13, 2024 · e0f0081 · e0f0081
1 parent 91b28ea
commit e0f0081
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/tests/test_hook.py b/tests/test_hook.py
@@ -45,6 +45,9 @@ def test_load():
     person = ab_2.Person(name='Bob')
     assert person == pickle.loads(pickle.dumps(person))
 
+    list_item = ab_2.container.ListItem()
+    assert list_item == pickle.loads(pickle.dumps(list_item))
+
 
 def test_load_module():
     ab = thriftpy2.load_module("addressbook_thrift")

diff --git a/thriftpy2/parser/__init__.py b/thriftpy2/parser/__init__.py
@@ -37,8 +37,16 @@ def load(path,
                    include_dir=include_dir, encoding=encoding)
     if incomplete_type:
         fill_incomplete_ttype(thrift, thrift)
+
+    # add sub modules to sys.modules recursively
     if real_module:
         sys.modules[module_name] = thrift
+        sub_modules = thrift.__thrift_meta__["includes"][:]
+        while sub_modules:
+            module = sub_modules.pop()
+            if module not in sys.modules:
+                sys.modules[module.__name__] = module
+                sub_modules.extend(module.__thrift_meta__["includes"])
     return thrift