Merge pull request #22 from tanmayaeron/master

resolved bugs and cleanup
descentis · Oct 3, 2022 · 44fe5b5 · 44fe5b5
2 parents aa0a916 + 56cece0
commit 44fe5b5
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 14 deletions.
diff --git a/kdap/analysis.py b/kdap/analysis.py
@@ -32,6 +32,7 @@
 from kdap.converter.wiki_clean import getCleanText
 from kdap.wiki_graph import graph_creater as gp
 from kdap.wikiextract.wikiExtract import wikiExtract
+from kdap.models.enums.granularity import Granularity
 
 
 class instances(object):
@@ -1475,16 +1476,16 @@ def __get_editor(self, *args, **kwargs):
                                                 if t > e:
                                                     editor_bool = 0
                                                     continue
-                                            if kwargs['granularity'].lower() == 'monthly':
+                                            if kwargs['granularity'].lower() == Granularity.MONTHLY.value:
                                                 if editor_dict.get(t.year) is None:
                                                     editor_dict[t.year] = {}
                                                     editor_dict[t.year][t.month] = []
                                                 elif editor_dict[t.year].get(t.month) is None:
                                                     editor_dict[t.year][t.month] = []
-                                            elif kwargs['granularity'].lower() == 'yearly':
+                                            elif kwargs['granularity'].lower() == Granularity.YEARLY.value:
                                                 if editor_dict.get(t.year) is None:
                                                     editor_dict[t.year] = []
-                                            elif kwargs['granularity'].lower() == 'daily':
+                                            elif kwargs['granularity'].lower() == Granularity.DAILY.value:
                                                 if editor_dict.get(t.year) is None:
                                                     editor_dict[t.year] = {}
                                                     editor_dict[t.year][t.month] = {}
@@ -1501,20 +1502,20 @@ def __get_editor(self, *args, **kwargs):
                                         U = chi.text
 
                                     if editor_bool:
+                                        if kwargs.get('granularity') is not None:
+                                            if kwargs['granularity'].lower() is not None:
+                                                if kwargs['granularity'].lower() == Granularity.MONTHLY.value:
 
-                                        if kwargs['granularity'].lower() is not None:
-                                            if kwargs['granularity'].lower() == 'monthly':
-
-                                                if U not in editor_dict[t.year][t.month]:
-                                                    editor_dict[t.year][t.month].append(U)
+                                                    if U not in editor_dict[t.year][t.month]:
+                                                        editor_dict[t.year][t.month].append(U)
 
-                                            elif kwargs['granularity'].lower() == 'daily':
-                                                if U not in editor_dict[t.year][t.month][t.day]:
-                                                    editor_dict[t.year][t.month][t.day].append(U)
+                                                elif kwargs['granularity'].lower() == Granularity.DAILY.value:
+                                                    if U not in editor_dict[t.year][t.month][t.day]:
+                                                        editor_dict[t.year][t.month][t.day].append(U)
 
-                                            elif kwargs['granularity'].lower() == 'yearly':
-                                                if U not in editor_dict[t.year]:
-                                                    editor_dict[t.year].append(U)
+                                                elif kwargs['granularity'].lower() == Granularity.YEARLY.value:
+                                                    if U not in editor_dict[t.year]:
+                                                        editor_dict[t.year].append(U)
                                     else:
                                         if U not in uList:
                                             uList.append(U)

diff --git a/kdap/models/__init__.py b/kdap/models/__init__.py
diff --git a/kdap/models/enums/__init__.py b/kdap/models/enums/__init__.py
diff --git a/kdap/models/enums/granularity.py b/kdap/models/enums/granularity.py
@@ -0,0 +1,7 @@
+from enum import Enum
+
+class Granularity(Enum):
+    """Enum for the granularity of the data."""
+    DAILY = 'daily'
+    MONTHLY = 'monthly'
+    YEARLY = 'yearly'