extract classmethods to try and fix tests

koxudaxi · Jun 25, 2024 · bdbfec7 · bdbfec7
1 parent 20b694e
commit bdbfec7
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 19 deletions.
diff --git a/datamodel_code_generator/parser/base.py b/datamodel_code_generator/parser/base.py
@@ -282,6 +282,18 @@ def process_module_tuple(input_tuple) -> Tuple[str, ...]:
     return tuple(r)
 
 
+def _postprocess_result_modules(results):
+    results = {process_module_tuple(k): v for k, v in results.items()}
+    init_result = [v for k, v in results.items() if k[-1] == '__init__.py'][0]
+    folders = {t[:-1] if t[-1].endswith('.py') else t for t in results.keys()}
+    for folder in folders:
+        for i in range(len(folder)):
+            subfolder = folder[: i + 1]
+            init_file = subfolder + ('__init__.py',)
+            results.update({init_file: init_result})
+    return results
+
+
 def _find_base_classes(model: DataModel) -> List[DataModel]:
     return [
         b.reference.source
@@ -1183,19 +1195,6 @@ def __set_one_literal_on_default(self, models: List[DataModel]) -> None:
                 if model_field.nullable is not True:  # pragma: no cover
                     model_field.nullable = False
 
-    @classmethod
-    def __postprocess_result_modules(cls, results):
-        results = {process_module_tuple(k): v for k, v in results.items()}
-
-        init_result = [v for k, v in results.items() if k[-1] == '__init__.py'][0]
-        folders = {t[:-1] if t[-1].endswith('.py') else t for t in results.keys()}
-        for folder in folders:
-            for i in range(len(folder)):
-                subfolder = folder[: i + 1]
-                init_file = subfolder + ('__init__.py',)
-                results.update({init_file: init_result})
-        return results
-
     def __change_imported_model_name(
         self,
         models: List[DataModel],
@@ -1375,7 +1374,7 @@ class Processed(NamedTuple):
 
         results = {tuple(i.replace('-', '_') for i in k): v for k, v in results.items()}
         results = (
-            self.__postprocess_result_modules(results)
+            _postprocess_result_modules(results)
             if self.treat_dots_as_module
             else {
                 tuple(

diff --git a/tests/parser/test_base.py b/tests/parser/test_base.py
@@ -8,8 +8,9 @@
 from datamodel_code_generator.parser.base import (
     Parser,
     exact_import,
+    process_module_tuple,
     relative,
-    sort_data_models, process_module_tuple,
+    sort_data_models,
 )
 from datamodel_code_generator.reference import Reference, snake_to_upper_camel
 from datamodel_code_generator.types import DataType
@@ -253,7 +254,7 @@ def test_no_additional_imports():
 
 
 def test_process_module_tuple():
-    t_1 = ("module1", "module2.module3", "module4.py")
-    t_2 = ("module1.module2", "module3.module4.py")
-    assert process_module_tuple(t_1) == ("module1", "module2", "module3", "module4.py")
-    assert process_module_tuple(t_2) == ("module1", "module2", "module3", "module4.py")
+    t_1 = ('module1', 'module2.module3', 'module4.py')
+    t_2 = ('module1.module2', 'module3.module4.py')
+    assert process_module_tuple(t_1) == ('module1', 'module2', 'module3', 'module4.py')
+    assert process_module_tuple(t_2) == ('module1', 'module2', 'module3', 'module4.py')