@krassowski/jupyterlab_go_to_definition/lib/languages/python.js

Jump to definition of a variable or function in JupyterLab

import { LanguageWithOptionalSemicolons } from './analyzer';
import { PathExt } from '@jupyterlab/coreutils';
function evaluateSkippingBrackets(tokens, indexShift, callback, allowNegativeBrackets = false) {
    // here `nextToken` is any token, not necessarily a meaningful one
    let nextToken = tokens[indexShift];
    // unpacking with curly braces is not possible
    const openingBrackets = '([';
    const closingBrackets = ')]';
    let openedBrackets = 0;
    // value of token is equal to an empty string for line breaks
    // a line break is the latest when the search should terminate,
    // unless the left-hand tuple is spread over several lines (in brackets)
    while (nextToken && (nextToken.value !== '' || openedBrackets > 0)) {
        if (nextToken.value === '') {
            // ignoring new-lines (when within brackets)
        }
        else if (openingBrackets.includes(nextToken.value)) {
            openedBrackets += 1;
        }
        else if (closingBrackets.includes(nextToken.value) &&
            (allowNegativeBrackets || openedBrackets > 0)) {
            openedBrackets -= 1;
        }
        else if (nextToken.value === ' ' || nextToken.value === '\t') {
            // ignoring whitespaces
        }
        else {
            let result = callback(nextToken, indexShift);
            if (result !== undefined) {
                return result;
            }
        }
        indexShift += 1;
        nextToken = tokens[indexShift];
    }
    return false;
}
const python_setup = `
import json
from pathlib import Path

def jupyter_lab_consumable_path(path):
    # if the definition is in current notebook, the path is not known on the kernel side
    if not path:
        return {'path': None, 'is_symlink': False}

    path = Path(path)
    cwd = Path.cwd()

    # the simple case
    if cwd in path.parents:
        return {
            'path': str(path.relative_to(cwd)),
            'is_symlink': False
        }

    symlinks_dir = Path('.jupyter_symlinks')
    symlinks_dir.mkdir(exist_ok=True)
    # relative pathways could lead out of the starting dir as well (e.g. using .. construct on Linux)
    path = path.absolute().resolve()

    # remove anchor (root slash/drive etc)
    sub_path = path.relative_to(path.anchor)

    # TODO: add drive to avoid confusion on Windows
    symlink = symlinks_dir / sub_path
    symlink.parent.mkdir(exist_ok=True, parents=True)
    try:
        symlink.unlink()
    except FileNotFoundError:
        pass
    assert not symlink.is_absolute()
    symlink.symlink_to(path)
    return {
        'path': str(symlink),
        'is_symlink': True
    }
`;
export class PythonAnalyzer extends LanguageWithOptionalSemicolons {
    constructor() {
        super(...arguments);
        // idea for improvement:
        //  rename Analyzer to RuleTester, define class Rule, make Rule instances take a callback on init,
        //  possibly add a string with rule's name (it could be displayed as "defined in >for< loop, line 9",
        //  or "imported from xyz module" and in case of multiple hits, user could choose which one to jump to),
        //  and make the rules interface the way to go for other languages.
        this.definitionRules = [
            this.isStandaloneAssignment,
            this.isImport,
            this.isWithStatement,
            this.isForLoopOrComprehension,
            this.isTupleUnpacking,
            this.isStoreMagic,
            this.isRMagicOutput
        ];
        this.supportsKernel = true;
    }
    // Matching standalone variable assignment:
    isStandaloneAssignment(siblings) {
        let { next } = siblings;
        return next.exists && this.isAssignment(next);
    }
    _is_magic_switch(candidate, key, max_args = 20) {
        while (max_args && candidate.exists) {
            if (candidate.value === key && candidate.simple_previous === '-') {
                break;
            }
            candidate = candidate.previous;
            max_args -= 1;
        }
        let is_switch = max_args !== 0;
        return {
            is_switch: is_switch,
            switch: is_switch ? candidate : null
        };
    }
    _is_magic_export(context, magic, export_arg, nargs = 1) {
        let { previous } = context;
        let switch_test = this._is_magic_switch(previous, export_arg, nargs);
        if (!switch_test.is_switch) {
            return false;
        }
        let magic_token = switch_test.switch.previous.previous;
        let percent = magic_token.simple_previous;
        return magic_token.value === magic && percent === '%';
    }
    // IPython %store -r magic:
    isStoreMagic(context) {
        return this._is_magic_export(context, 'store', 'r', 20);
    }
    isRMagicOutput(context) {
        return this._is_magic_export(context, 'R', 'o', 1);
    }
    // Matching imports:
    isImport(context) {
        let { previous } = context;
        return (previous.exists &&
            previous.type === 'keyword' &&
            previous.value === 'import');
    }
    isCrossFileReference(context) {
        // from a import b; from a.b import c
        let previous = context.previous;
        let next = context.next;
        previous = this.traverse_left(previous, '.');
        next = this.traverse_right(next, '.');
        if (previous.exists &&
            previous.type === 'keyword' &&
            previous.value === 'from' &&
            next.exists &&
            next.type === 'keyword' &&
            next.value === 'import') {
            return true;
        }
        // import x, import a.b
        let before_previous = previous.previous.previous;
        if (this.isImport(previous.next) &&
            !(before_previous.exists &&
                before_previous.type === 'keyword' &&
                before_previous.value === 'from')) {
            return true;
        }
        return false;
    }
    _breadcrumbs(context, allow_import_dots_upfront = false) {
        let { previous, token } = context;
        let parts = [];
        let is_dot = previous.simple_next === '.';
        while (is_dot && previous.exists) {
            parts.push(previous.value);
            previous = previous.previous;
            is_dot = previous.simple_next === '.';
        }
        if (allow_import_dots_upfront) {
            // relative imports
            if (previous.simple_previous === '.') {
                parts.push('');
            }
            if (previous.simple_previous === '..') {
                parts.push('.');
            }
        }
        parts = parts.reverse();
        parts.push(token.value);
        return parts;
    }
    _imports_breadcrumbs(context) {
        return this._breadcrumbs(context, true);
    }
    definitionLocationQuery(context) {
        let parts = this._breadcrumbs(context);
        let value = parts.join('.');
        if (/^[a-zA-Z_.]+$/.test(value)) {
            return (python_setup +
                `

def _extract_object(breadcrumbs, available_objects):
    obj = available_objects[breadcrumbs[0]]

    for name in breadcrumbs[1:]:
        obj = getattr(obj, name)
    return obj


def _locate_definition(name):
    """Returns (line_number, path, is_sym_link) tuple"""
    from inspect import getsourcefile, getsourcelines

    available_objects = globals()
    input_number = None

    try:
        obj = _extract_object(name.split('.'), available_objects)
    except (KeyError, AttributeError):
        return

    try:
        path = getsourcefile(obj)
        path_exists = Path(path).exists()
        if not path_exists:
            assert path.startswith('<ipython-input-')
            input_number = int(path.split('-')[2])
            path = None
    except (TypeError, AssertionError):
        path = None

    try:
        source, line_number = getsourcelines(obj)
    except OSError:
        line_number = 0

    return {
        'line_number': line_number,
        'input_number': input_number,
        'column': len(source[0]) - len(source[0].lstrip()),
        **jupyter_lab_consumable_path(path)
    }

print(json.dumps(_locate_definition('` +
                value +
                `')), end='')
`);
        }
    }
    referencePathQuery(context) {
        let parts = this._imports_breadcrumbs(context);
        let value = parts.join('.');
        // TODO: recognize Python version and return no query (or another Python2-compatible query)
        if (/^[a-zA-Z_.]+$/.test(value)) {
            // just in case to prevent arbitrary execution
            return (python_setup +
                `
def _get_path(value):
    """Returns (path, is_sym_link) tuple"""
    from importlib.util import find_spec

    path = find_spec(value).origin

    return jupyter_lab_consumable_path(path)

print(json.dumps(_get_path('` +
                value +
                `')), end='')
`);
        }
    }
    guessReferencePath(context) {
        let parts = this._imports_breadcrumbs(context);
        let prefix = parts.join('/');
        return [prefix + '.py', PathExt.join(prefix, '__init__.py')];
    }
    // Matching `as`:
    // e.g. `with open('name') as f:` or `except Exception as e:`
    isWithStatement(context) {
        let { previous } = context;
        return (previous.exists && previous.type === 'keyword' && previous.value === 'as');
    }
    // Matching `for` loop and comprehensions:
    isForLoopOrComprehension(context) {
        let { previous, next } = context;
        return (previous.exists &&
            previous.type === 'keyword' &&
            previous.value === 'for' &&
            next.exists &&
            next.type === 'keyword' &&
            next.value === 'in');
    }
    isTupleUnpacking(context) {
        // Matching variables in tuple unpacking:
        let { tokens, index } = context;
        // Considering: `a, [b, c], (d, ) = 1, [1, 2], (1,)`, if the tested
        // token is `a`, then the next expected token would be a comma,
        // and then one of following: a variable, an assignment symbol,
        // or an opening bracket (for simplicity brackets can be ignored).
        let commaExpected = true;
        return evaluateSkippingBrackets(tokens, index + 1, (nextToken, indexShift) => {
            if (nextToken.type === 'operator' && nextToken.value === '=') {
                let lastToken;
                evaluateSkippingBrackets(tokens, indexShift + 1, (nextToken) => {
                    lastToken = nextToken;
                });
                // return true unless in a function call
                if (!lastToken || lastToken.value !== ')') {
                    return true;
                }
            }
            if (commaExpected && nextToken.value !== ',') {
                return false;
            }
            commaExpected = !commaExpected;
        }, true);
    }
}
//# sourceMappingURL=python.js.map