From 3597481fe88f6538067c0ae5c252fd6fdd3f8b81 Mon Sep 17 00:00:00 2001
From: Rangeet Pan <Rangeet.Pan@ibm.com>
Date: Mon, 9 Sep 2024 10:08:39 -0400
Subject: [PATCH] update for python class name

---
 cldk/analysis/python/treesitter/python_sitter.py | 1 +
 cldk/models/python/models.py                     | 1 +
 2 files changed, 2 insertions(+)

diff --git a/cldk/analysis/python/treesitter/python_sitter.py b/cldk/analysis/python/treesitter/python_sitter.py
index ec3ba656..df5dcc3f 100644
--- a/cldk/analysis/python/treesitter/python_sitter.py
+++ b/cldk/analysis/python/treesitter/python_sitter.py
@@ -197,6 +197,7 @@ def get_all_classes(self, module: str) -> List[PyClass]:
                                    full_signature=class_full_signature,
                                    methods=methods,
                                    super_classes=super_classes,
+                                   class_name=klass_name,
                                    is_test_class=is_test_class))
         return classes
 
diff --git a/cldk/models/python/models.py b/cldk/models/python/models.py
index ab08b94c..1f398f8a 100644
--- a/cldk/models/python/models.py
+++ b/cldk/models/python/models.py
@@ -45,6 +45,7 @@ class PyClass(BaseModel):
     full_signature: str
     super_classes: List[str]
     is_test_class: bool
+    class_name: str = None
     methods: List[PyMethod]