Merge pull request #1 from FSoft-AI4Code/feature/return_type

Feature/return type

Author: minhna1112

src/codetext/parser/c_sharp_parser.py

@@ -105,15 +105,13 @@ def get_class_list(node):
         return res
 
     @staticmethod
-    def get_function_metadata(function_node, blob: str=None) -> Dict[str, Any]:
+    def get_function_metadata(function_node, blob: str = None) -> Dict[str, Any]:
         """
         Function metadata contains:
             - identifier (str): function name
             - parameters (Dict[str, str]): parameter's name and their type (e.g: {'param_a': 'int'})
             - type (str): type
         """
-        if blob:
-            logger.info('From version `0.0.6` this function will update argument in the API')
         metadata = {
             'identifier': '',
             'parameters': {},
@@ -122,10 +120,13 @@ def get_function_metadata(function_node, blob: str=None) -> Dict[str, Any]:
         assert type(function_node) == tree_sitter.Node
 
         for child in function_node.children:
-            if child.type == 'predefined_type':
+            if child.type in ['predefined_type', 'generic_name']:
                 metadata['return_type'] = get_node_text(child)
             elif child.type == 'identifier':
-                metadata['identifier'] = get_node_text(child)
+                if child.next_named_sibling.type != 'parameter_list':
+                    metadata['return_type'] = get_node_text(child)
+                else:
+                    metadata['identifier'] = get_node_text(child)
             elif child.type == 'parameter_list':
                 for param_node in child.children:
                     param_nodes = get_node_by_kind(param_node, ['parameter'])

src/codetext/parser/java_parser.py

@@ -103,19 +103,26 @@ def get_class_metadata(class_node, blob: str=None) -> Dict[str, str]:
         return metadata
 
     @staticmethod
-    def get_function_metadata(function_node, blob: str=None) -> Dict[str, str]:
-        if blob:
-            logger.info('From version `0.0.6` this function will update argument in the API')
+    def get_function_metadata(function_node, blob: str = None) -> Dict[str, str]:
         metadata = {
             'identifier': '',
             'parameters': {},
             'return_type': None
         }
 
+        return_kinds = ["void_type", 
+                        "integral_type",
+                        "floating_point_type",
+                        "boolean_type",
+                        "type_identifier",
+                        "scoped_type_identifier",
+                        "generic_type"]
+        
+
         for child in function_node.children:
             if child.type == 'identifier':
-                metadata['identifier'] = get_node_text(child)
-            elif child.type == 'type_identifier':
+                metadata['identifier'] = get_node_text(child)    
+            elif child.type in return_kinds:
                 metadata['return_type'] = get_node_text(child)
             elif child.type == 'formal_parameters':
                 param_list = get_node_by_kind(child, ['formal_parameter'])  # speed_parameter
@@ -124,4 +131,5 @@ def get_function_metadata(function_node, blob: str=None) -> Dict[str, str]:
                     identifier = get_node_text(param.child_by_field_name('name'))
                     metadata['parameters'][identifier] = param_type
 
-        return metadata
+        
+        return metadata