Merge pull request emrgnt-cmplxty#146 from maks-ivanov/feature/polish-search

Feature/polish search

Author: emrgnt-cmplxty

automata/tools/search/local_types.py

@@ -1,7 +1,8 @@
 import re
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, Dict, List, Optional
+from os import PathLike
+from typing import Any, Dict, List, Optional, Union
 
 from automata.tools.search.scip_pb2 import Descriptor as DescriptorProto
 
@@ -161,6 +162,7 @@ def __eq__(self, other):
 
 @dataclass
 class SymbolReference:
+    symbol: Symbol
     line_number: int
     roles: Dict[str, Any]
 
@@ -179,3 +181,6 @@ def __eq__(self, other):
         elif isinstance(other, str):
             return self.path == other
         return False
+
+
+StrPath = Union[str, PathLike]

automata/tools/search/main.py

@@ -1,7 +1,9 @@
 import os
 from argparse import ArgumentParser
+from typing import Dict, cast
 
 # from automata.tools.search.call_graph import CallGraph
+# from automata.tools.search.local_types import Descriptor
 from automata.tools.search.symbol_converter import SymbolConverter
 from automata.tools.search.symbol_graph import SymbolGraph
 from automata.tools.search.symbol_parser import parse_uri_to_symbol
@@ -25,17 +27,17 @@
     # Dump all available files in the symbol graph
     print("-" * 200)
     print("Fetching all available files in SymbolGraph")
-    file_nodes = symbol_graph.get_all_files()
-    for file_node in file_nodes:
-        print("File >> %s" % (file_node))
+    files = symbol_graph.get_all_files()
+    for file in files:
+        print("File Path >> %s" % (file.path))
     print("-" * 200)
 
-    # Dump all available symbols at the test path
+    # Dump all available symbols defined along test path
     print("-" * 200)
-    print("Fetching all available symbols along %s" % (test_path))
+    print("Fetching all defined symbols along %s" % (test_path))
     available_symbols = symbol_graph.get_defined_symbols_along_path(test_path)
     for symbol in available_symbols:
-        print("Available Symbol >> %s" % (symbol))
+        print("Defined Symbol >> %s" % (symbol))
 
     print("-" * 200)
 
@@ -48,23 +50,30 @@
     # Find references of the test symbol
     print("-" * 200)
     print("Searching for references of the symbol %s" % (test_symbol))
-    search_result = symbol_searcher.process_query("type:symbol %s" % (test_symbol.uri))
-    print("References: ", search_result)
+    search_result_0: Dict = cast(
+        Dict, symbol_searcher.process_query("type:symbol %s" % (test_symbol.uri))
+    )
+    for file_path in search_result_0.keys():
+        print("File Path >> %s" % (file_path))
+        for reference in search_result_0[file_path]:
+            print("Reference >> %s" % (reference))
     print("-" * 200)
 
     # Find source code for the test symbol
     print("-" * 200)
     print("Searching for source code for symbol %s" % (test_symbol))
-    search_result = symbol_searcher.process_query("type:source %s" % (test_symbol.uri))
-    print("Source Code: ", search_result)
+    search_result_1: str = cast(
+        str, symbol_searcher.process_query("type:source %s" % (test_symbol.uri))
+    )
+    print("Source Code: ", search_result_1)
     print("-" * 200)
 
     # Find exact matches for abbrievated test symbol
     print("-" * 200)
     abbv_test_symbol = "AutomataAgentConfig"
     print("Searching for exact matches of the filter %s" % (abbv_test_symbol))
-    search_result = symbol_searcher.process_query("type:exact %s" % (abbv_test_symbol))
-    print("Search result: ", search_result)
+    search_result_2 = symbol_searcher.process_query("type:exact %s" % (abbv_test_symbol))
+    print("Search result: ", search_result_2)
     print("-" * 200)
 
     # Perform a find and replace on the test find symbol below

automata/tools/search/symbol_graph.py

@@ -3,7 +3,7 @@
 
 import networkx as nx
 
-from automata.tools.search.local_types import File, Symbol, SymbolReference
+from automata.tools.search.local_types import File, StrPath, Symbol, SymbolReference
 from automata.tools.search.scip_pb2 import Index, SymbolRole
 from automata.tools.search.symbol_converter import SymbolConverter
 from automata.tools.search.symbol_parser import parse_uri_to_symbol
@@ -13,7 +13,10 @@
 
 class SymbolGraph:
     def __init__(
-        self, index_path: str, symbol_converter: SymbolConverter, do_shortened_symbols: bool = True
+        self,
+        index_path: StrPath,
+        symbol_converter: SymbolConverter,
+        do_shortened_symbols: bool = True,
     ):
         """
         Initialize SymbolGraph with the path of an index protobuf file.
@@ -33,7 +36,11 @@ def get_all_files(self) -> List[File]:
 
         :return: List of all file nodes.
         """
-        return [data for _, data in self._graph.nodes(data=True) if data.get("label") == "file"]
+        return [
+            data.get("file")
+            for _, data in self._graph.nodes(data=True)
+            if data.get("label") == "file"
+        ]
 
     def get_all_defined_symbols(self) -> List[Symbol]:
         """
@@ -54,9 +61,9 @@ def get_symbol_references(self, symbol: Symbol) -> Dict[str, List[SymbolReferenc
         :return: List of tuples (file, symbol details)
         """
         search_results = [
-            (file_path, symbol_reference)
-            for file_path, symbol_reference, label in self._graph.out_edges(symbol, data=True)
-            if label == "reference"
+            (file_path, data.get("symbol_reference"))
+            for file_path, _, data in self._graph.out_edges(symbol, data=True)
+            if data.get("label") == "reference"
         ]
         result_dict: Dict[str, List[SymbolReference]] = {}
 
@@ -141,11 +148,12 @@ def _build_symbol_graph(self, index: Index) -> nx.MultiDiGraph:
         G = nx.MultiDiGraph()
         for document in index.documents:
             # Add File Vertices
-            document_path = document.relative_path
+            document_path: StrPath = document.relative_path
+
             G.add_node(
                 document_path,
                 file=File(path=document.relative_path, occurrences=document.occurrences),
-                label="file_path",
+                label="file",
             )
 
             for symbol_information in document.symbols:
@@ -167,6 +175,7 @@ def _build_symbol_graph(self, index: Index) -> nx.MultiDiGraph:
                 occurrence_range = tuple(occurrence_information.range)
                 occurrence_roles = self._get_symbol_roles_dict(occurrence_information.symbol_roles)
                 occurrence_reference = SymbolReference(
+                    symbol=occurrence_symbol,
                     line_number=occurrence_range[0],
                     roles=occurrence_roles,
                 )
@@ -195,7 +204,7 @@ def _get_symbol_roles_dict(role) -> Dict[str, bool]:
         return result
 
     @staticmethod
-    def _load_index_protobuf(path: str) -> Index:
+    def _load_index_protobuf(path: StrPath) -> Index:
         """
         Load an index from a protobuf file.