diff --git a/src/umlizer/class_graph.py b/src/umlizer/class_graph.py index b7a9340..4b54eb8 100644 --- a/src/umlizer/class_graph.py +++ b/src/umlizer/class_graph.py @@ -309,7 +309,11 @@ def _extract_module_name(module_path: str) -> tuple[str, str]: module_split = module_path.split(os.sep) module_path = os.sep.join(module_split[:-1]) module_filename = module_split[-1] - module_name = module_filename.rstrip('.py') + + if module_filename.endswith('.py'): + module_name = module_filename[:-3] + else: + module_name = module_filename return module_path, module_name diff --git a/src/umlizer/cli.py b/src/umlizer/cli.py index 8e3c72d..0f4b375 100644 --- a/src/umlizer/cli.py +++ b/src/umlizer/cli.py @@ -109,6 +109,12 @@ def class_( ) ), ] = '', + django_settings: Annotated[ + str, + typer.Option( + help='Django settings module (eg. "config.settings.dev").' + ), + ] = '', verbose: Annotated[ bool, typer.Option(help='Active the verbose mode.') ] = False, @@ -117,6 +123,11 @@ def class_( source = make_absolute(source) target = make_absolute(target) / 'class_graph' + if django_settings: + from umlizer.plugins import django + + django.setup(django_settings) + classes_nodes = class_graph.load_classes_definition( source, exclude=exclude, verbose=verbose ) diff --git a/src/umlizer/plugins/__init__.py b/src/umlizer/plugins/__init__.py new file mode 100644 index 0000000..44d7d43 --- /dev/null +++ b/src/umlizer/plugins/__init__.py @@ -0,0 +1 @@ +"""Set of functions for integrating to another libraries.""" diff --git a/src/umlizer/plugins/django.py b/src/umlizer/plugins/django.py new file mode 100644 index 0000000..a1140f6 --- /dev/null +++ b/src/umlizer/plugins/django.py @@ -0,0 +1,18 @@ +"""Set of functions for integrating to django.""" + +import os + + +def setup(settings_module: str) -> None: + """ + Set up the Django environment. + + Parameters + ---------- + settings_module : str + The Django settings module to use. + """ + import django + + os.environ['DJANGO_SETTINGS_MODULE'] = settings_module + django.setup()