Parsing/navigating through the Python AST with talon

pyast.py

import ast
import os.path

from talon import Context, Module, actions, ui

ctx = Context()
mod = Module()

# Stores the original AST elements so we can refer to specific properties of them
REFERENCE = {}

mod.list("ast_functions_current_file", desc="Functions in the current file")
mod.list("ast_variables_current_file", desc="Variables in the current file")

@mod.action_class
class Actions:
    """Logic for parsing the AST a files and generating lists"""

    def walk_ast(source: str):
        """Iterates through the source code and returns a dictionary"""
        root = ast.parse(source)

        result = {'variables': {}, 'functions': {}}

        for node in ast.walk(root):
            if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Store):
                result["variables"][node.id] = node
            if isinstance(node, ast.FunctionDef):
                result["functions"][node.name] = node

        return result

    def walk_and_set_lists(path: str):
        """Walks the AST of the given path and updates our Talon lists"""
        global REFERENCE

        if not path.endswith(".py") or not os.path.exists(path):
            return

        with open(path, "r") as f:
            result = actions.user.walk_ast(f.read())
            for kind, dict in result.items():
                kind_list = actions.user.create_spoken_forms_from_list(dict.keys())
                ctx.lists[f"self.ast_{kind}_current_file"] = kind_list
            REFERENCE = result

        actions.user.pretty_print(REFERENCE)

    def navigate_to_function(function_name: str):
        """Moves the editor to the function by the given name"""
        print(f"navigate to {function_name}")
        el = REFERENCE["functions"][function_name]
        actions.edit.jump_line(el.lineno)
        actions.edit.jump_column(el.end_col_offset)


ui.register("win_title", lambda x: actions.user.walk_and_set_lists(actions.user.file_manager_current_path()))
ui.register("win_focus", lambda x: actions.user.walk_and_set_lists(actions.user.file_manager_current_path()))
ui.register("app_focus", lambda x: actions.user.walk_and_set_lists(actions.user.file_manager_current_path()))

pyast.talon

(go to|function) {user.ast_functions_current_file}$: user.navigate_to_function(ast_functions_current_file)