core: added get_all_class_method_bodies for getting class method bodies just like functions

Author: rcshubhadeep

setup.py

@@ -23,7 +23,7 @@
         "Programming Language :: Python :: 3",
         "Programming Language :: Python :: 3.6",
 ],
-version="0.9.2",
+version="0.10.0",
 packages=find_packages(exclude=("tests",)),
 install_requires=["tree-sitter", "pygit2", "pytest", "PyYAML"],
 entry_points = {

tests/assets/file_with_different_functions.py

@@ -91,3 +91,6 @@ class AnotherClass(BaseClass):
     def __init__(self):
         super(AnotherClass, self).__init__(12)
         self.name = "AnotherClass"
+    
+    def get_class_name(self):
+        return self.name

tests/test_python_parser.py

@@ -66,7 +66,8 @@ def test_parser_get_all_class_method_names(python_parser):
 			'scan'
 		], 
 		'AnotherClass': [
-			'__init__'
+			'__init__',
+			'get_class_name'
 		]
 	}
 
@@ -86,6 +87,22 @@ def test_parser_get_all_function_bodies(python_parser):
 	)
 	assert python_parser.get_all_function_bodies()["function_different_args"] == \
 		'def function_different_args(foo, bar, no, pi, opt_typed):\n    return 0'
+
+def test_parser_get_all_method_bodies(python_parser):
+	assert python_parser.get_all_class_method_bodies()['BaseClass']["get_name"] == \
+		"def get_name(self):\n        '''Get the name parameter\n        '''\n        return self.name"
+	assert python_parser.get_all_class_method_bodies(get_index=True)['BaseClass']["get_name"] ==(
+		"def get_name(self):\n        '''Get the name parameter\n        '''\n        return self.name",
+   		(75, 8),
+   		(77, 24)
+	)
+	assert python_parser.get_all_class_method_bodies(get_index=True, strip_docstr=True)['BaseClass']["get_name"] == (
+		('def get_name(self):\n        return self.name',
+    	"'''Get the name parameter\n        '''"),
+   		(75, 8),
+   		(77, 24)
+	)
+
 
 def test_parser_get_all_function_documentations(python_parser):
 	assert python_parser.get_all_function_documentations() == {
@@ -98,11 +115,11 @@ def test_parser_get_all_function_documentations(python_parser):
 def test_parser_walk(python_parser):
 	l = []
 	python_parser.walk(lambda node: l.append(1) if node.type == "function_definition" else l.append(0))
-	assert sum(l) == 12
+	assert sum(l) == 13
 
 def test_parser_reduction(python_parser):
 	assert python_parser.reduction(
 	    lambda node,acc: acc+1 if node.type == "function_definition" else acc, 0
-	) == 12
+	) == 13
 
 

tree_hugger/core/parser/python/python_parser.py

@@ -64,7 +64,7 @@ def _outer_indent(self, code):
         return min(spaces_arr) if spaces_arr else 4
 
 
-    def get_all_class_method_names(self) -> List:
+    def get_all_class_method_names(self) -> Dict:
         """
         Gets all the method names from a file. 
 
@@ -164,6 +164,63 @@ def get_all_method_documentations(self, strip_quotes: bool=False) -> Dict:
     def _get_var_names_as_string_from_param_list(self, param_list):
         return "(" + ", ".join([var_name for var_name, _, _ in param_list]) + ")"
 
+    def get_all_class_method_bodies(self, strip_docstr: bool=False, get_index: bool=False) -> Dict:
+        class_method_names = self.get_all_class_method_names()
+        methods_and_params = self.get_all_function_names_with_params()
+        method_and_docstr = self.get_all_method_docstrings()
+
+        captures = self._run_query_and_get_captures('all_function_bodies', self.root_node)
+
+        ret_struct = {}
+        pp = {}
+
+        for i in range(0, len(captures), 2):
+            func_name = match_from_span(captures[i][0], self.splitted_code)
+            for class_name, method_name_list in class_method_names.items():
+                if func_name in method_name_list and (func_name != "__init__" and func_name != "__new__"):
+                    if pp.get(class_name) is None:
+                        pp[class_name] = {}
+                        pp[class_name] = {func_name: 
+                                          (match_from_span(captures[i+1][0], self.splitted_code), captures[i+1][0].start_point, captures[i+1][0].end_point)
+                                          }
+                    else:
+                        if func_name != "__init__" and func_name != "__new__":
+                            pp[class_name][func_name] =  (match_from_span(captures[i+1][0], self.splitted_code), captures[i+1][0].start_point, captures[i+1][0].end_point)
+
+        if strip_docstr:
+            for class_name, func_data_struct in pp.items():
+                ret_struct[class_name] = {}
+                for k, (v,sp,ep) in func_data_struct.items():
+                    if method_and_docstr[class_name].get(k) is not None and method_and_docstr[class_name].get(k) is not '':
+                        code = v.replace(method_and_docstr[class_name][k], "")
+                        outer_indent = self._outer_indent(code)
+                        spaces = " ".join([''] * (outer_indent + 1))
+                        if code.startswith("\n"):
+                            ret_struct[class_name][k] = (f"def {k}{self._get_var_names_as_string_from_param_list(methods_and_params[k])}:{code}", 
+                                                         method_and_docstr[class_name][k])
+                        else:
+                            ret_struct[class_name][k] = (f"def {k}{self._get_var_names_as_string_from_param_list(methods_and_params[k])}:\n{spaces}{code}",
+                                                         method_and_docstr[class_name][k])
+                    else:
+                        outer_indent = self._outer_indent(v)
+                        spaces = " ".join([''] * (outer_indent + 1))
+                        ret_struct[class_name][k] = (f"def {k}{self._get_var_names_as_string_from_param_list(methods_and_params[k])}:\n{spaces}{v}", "")
+                    if get_index:
+                        ret_struct[class_name][k] = ret_struct[class_name][k],sp,ep
+        else:
+            for class_name, func_data_struct in pp.items():
+                for k, (v,sp,ep) in func_data_struct.items():
+                    outer_indent = self._outer_indent(v)
+                    spaces = " ".join([''] * (outer_indent + 1))
+                    if ret_struct.get(class_name) is None:
+                        ret_struct[class_name] = {}
+                        ret_struct[class_name][k] = f"def {k}{self._get_var_names_as_string_from_param_list(methods_and_params[k])}:\n{spaces}{v}"
+                    else:
+                        ret_struct[class_name][k] = f"def {k}{self._get_var_names_as_string_from_param_list(methods_and_params[k])}:\n{spaces}{v}"
+                    if get_index:
+                        ret_struct[class_name][k] = (ret_struct[class_name][k], sp, ep)
+        return ret_struct
+    
     def get_all_function_bodies(self, strip_docstr: bool=False, get_index: bool=False) -> Dict:
         """
         Returns a dict where function names are the key and the whole function code are the values