Function touch_import

tools/python-3.11.9-amd64/Lib/lib2to3/fixer_util.py:315–357

Works like `does_tree_import` but adds an import statement if it was not imported.

(package, name, node)

    return node.type in (syms.import_name, syms.import_from)

def touch_import(package, name, node):
    """ Works like `does_tree_import` but adds an import statement
        if it was not imported. """
    def is_import_stmt(node):
        return (node.type == syms.simple_stmt and node.children and
                is_import(node.children[0]))

    root = find_root(node)

    if does_tree_import(package, name, root):
        return

    # figure out where to insert the new import.  First try to find
    # the first import and then skip to the last one.
    insert_pos = offset = 0
    for idx, node in enumerate(root.children):
        if not is_import_stmt(node):
            continue
        for offset, node2 in enumerate(root.children[idx:]):
            if not is_import_stmt(node2):
                break
        insert_pos = idx + offset
        break

    # if there are no imports where we can insert, find the docstring.
    # if that also fails, we stick to the beginning of the file
    if insert_pos == 0:
        for idx, node in enumerate(root.children):
            if (node.type == syms.simple_stmt and node.children and
               node.children[0].type == token.STRING):
                insert_pos = idx + 1
                break

    if package is None:
        import_ = Node(syms.import_name, [
            Leaf(token.NAME, "import"),
            Leaf(token.NAME, name, prefix=" ")
        ])
    else:
        import_ = FromImport(package, [Leaf(token.NAME, name, prefix=" ")])

    children = [import_, Newline()]
    root.insert_child(insert_pos, Node(syms.simple_stmt, children))


_def_syms = {syms.classdef, syms.funcdef}