44'''
55
66import types
7+ from typing import overload
78
89from flatdata .generator .tree .builder import build_ast
910from flatdata .generator .tree .nodes .trivial .namespace import Namespace
1011from flatdata .generator .tree .nodes .node import Node
12+ from flatdata .generator .tree .syntax_tree import SyntaxTree
1113
1214from .generators .cpp import CppGenerator
1315from .generators .dot import DotGenerator
1416from .generators .go import GoGenerator
1517from .generators .python import PythonGenerator
1618from .generators .rust import RustGenerator
1719from .generators .flatdata import FlatdataGenerator
20+ from .generators import BaseGenerator
1821
1922
2023class Engine :
@@ -23,7 +26,7 @@ class Engine:
2326 Implements code generation from the given flatdata schema.
2427 """
2528
26- _GENERATORS = {
29+ _GENERATORS : dict [ str , type [ BaseGenerator ]] = {
2730 "cpp" : CppGenerator ,
2831 "dot" : DotGenerator ,
2932 "go" : GoGenerator ,
@@ -33,21 +36,21 @@ class Engine:
3336 }
3437
3538 @classmethod
36- def available_generators (cls ):
39+ def available_generators (cls ) -> list [ str ] :
3740 """
3841 Lists names of available code generators.
3942 """
4043 return list (cls ._GENERATORS .keys ())
4144
42- def __init__ (self , schema ) :
45+ def __init__ (self , schema : str ) -> None :
4346 """
4447 Instantiates generator engine for a given schema.
4548 :raises FlatdataSyntaxError
4649 """
4750 self .schema = schema
4851 self .tree = build_ast (schema )
4952
50- def render (self , generator_name ) :
53+ def render (self , generator_name : str ) -> str :
5154 """
5255 Render schema with a given generator
5356 :param generator_name:
@@ -60,38 +63,45 @@ def render(self, generator_name):
6063 )
6164
6265 output_content = generator .render (self .tree )
63- return output_content
66+ return str ( output_content )
6467
65- def render_python_module (self , module_name = None , archive_name = None , root_namespace = None ):
68+ @overload
69+ def render_python_module (self , module_name : str | None , archive_name : str , root_namespace : str | None = None ) -> tuple [types .ModuleType , type ]: ...
70+ @overload
71+ def render_python_module (self , * , archive_name : str , root_namespace : str | None = None ) -> tuple [types .ModuleType , type ]: ...
72+ @overload
73+ def render_python_module (self , module_name : str | None = None , archive_name : None = None , root_namespace : str | None = None ) -> types .ModuleType : ...
74+
75+ def render_python_module (self , module_name : str | None = None , archive_name : str | None = None , root_namespace : str | None = None ) -> types .ModuleType | tuple [types .ModuleType , type ]:
6676 """
6777 Render python module.
6878 :param module_name: Module name to use. If none, root namespace name is used.
6979 :param archive_name: Archive name to lookup,
7080 if specified, archive type is returned along with the model
7181 :param root_namespace: Root namespace to pick in case of multiple top level namespaces.
7282 """
73- root_namespace = self ._find_root_namespace (self .tree , archive_name , root_namespace )
83+ ns = self ._find_root_namespace (self .tree , archive_name , root_namespace )
7484 module_code = self .render ("py" )
75- module = types .ModuleType (module_name if module_name is not None else root_namespace .name )
85+ module = types .ModuleType (module_name if module_name is not None else ns .name )
7686 #pylint: disable=exec-used
7787 exec (module_code , module .__dict__ )
7888 if archive_name is None :
7989 return module
8090
81- name = root_namespace .name + "_" + archive_name
82- archive_type = getattr (module , name ) if archive_name else None
91+ name = ns .name + "_" + archive_name
92+ archive_type = getattr (module , name )
8393 return module , archive_type
8494
8595 @classmethod
86- def _create_generator (cls , name ) :
96+ def _create_generator (cls , name : str ) -> BaseGenerator | None :
8797 generator_type = cls ._GENERATORS .get (name , None )
8898 if generator_type is None :
8999 return None
90100
91- return generator_type ()
101+ return generator_type () # type: ignore[call-arg] # dict values are concrete subclasses with zero-arg __init__
92102
93103 @staticmethod
94- def _find_root_namespace (tree , archive_name , root_namespace = None ):
104+ def _find_root_namespace (tree : SyntaxTree , archive_name : str | None , root_namespace : str | None = None ) -> Namespace :
95105 root_children = tree .root .children
96106 root_namespaces = [
97107 child for child in root_children
0 commit comments