Recipe for adding an import

[zhammer]

Curious if you have a preferred way of adding an import to a module. This was a quick solution I put together, but it wouldn't work in a lot of scenarios and is messy:

    @staticmethod
    def _with_added_import(module_node: cst.Module, import_node: cst.Import) -> cst.Module:
        """
        Adds new import `import_node` after the first import in the module `module_node`.
        """
        updated_body: List[Union[cst.SimpleStatementLine, BaseCompoundStatement]] = []
        added_import = False
        for line in module_node.body:
            updated_body.append(line)
            if not added_import and IterItemsTransformer._is_import_line(line):
                updated_body.append(cst.SimpleStatementLine(body=tuple([import_node])))
                added_import = True

        return module_node.with_changes(body=tuple(updated_body))


    @staticmethod
    def _is_import_line(line: Union[cst.SimpleStatementLine, BaseCompoundStatement]) -> bool:
        return (
            isinstance(line, cst.SimpleStatementLine) and
            len(line.body) == 1 and
            isinstance(line.body[0], (cst.Import, cst.ImportFrom))
        )

[DragonMinded]

We have an additional layer that we use internally which includes an add_needed_import transform. I want to see it open-sourced (we're effectively using our internal repository as an incubator for all of this at this point) but we aren't satisfied with the code quality yet.

Effectively we are using a separate visitor to gather up all of the import nodes, and then calling visit() on the tree inside visit_Module. Note that across multiple visitors (not transforms, since they return a new tree) nodes are comparable via identity so you can stick nodes you care about in a list/set/dict and then do x in y checks against it later. This gives us all of the existing imports we care about, which lets us add to an existing import or create a new one using code similar to what you've posted. Part of the reason this didn't make open-source is a lot of the visitor behavior for gathering up existing nodes will be overcome by scope analysis work that @jimmylai is working on. This should be coming in the next few weeks and will make its way into libcst.metadata.

Long term, one of the things I'd like to see in this repository is a library of transforms that do high-level actions such as adding/removing imports. Sort of a batteries-included set of actions which abstract automated modifications to the next level. I'm super excited about all of that but there's a ton of work that we need to get to before it becomes a reality.

[zhammer]

this is super helpful, thank you @DragonMinded. once your team has the bandwidth, it could be great to have a spectrum channel (or some other place) for libcst questions/conversations.

[DragonMinded]

For now, we're totally okay with opening issues for questions. Keep 'em coming!