From 9df3c54ecf01505323421499d7e8aecb14bdd145 Mon Sep 17 00:00:00 2001
From: Ivan Ogasawara <ivan.ogasawara@gmail.com>
Date: Tue, 25 Jun 2024 21:50:23 -0400
Subject: [PATCH] fix issues with rstrip

---
 src/umlizer/class_graph.py      |  6 +++++-
 src/umlizer/cli.py              | 11 +++++++++++
 src/umlizer/plugins/__init__.py |  1 +
 src/umlizer/plugins/django.py   | 18 ++++++++++++++++++
 4 files changed, 35 insertions(+), 1 deletion(-)
 create mode 100644 src/umlizer/plugins/__init__.py
 create mode 100644 src/umlizer/plugins/django.py

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()