from typing import Union

import libcst as cst


class RemoveUnusedImportTransformer(cst.CSTTransformer):
    def leave_ImportFrom(
        self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom
    ) -> Union[cst.RemovalSentinel, cst.ImportFrom]:
        if updated_node.names[0].name.value == "token":  # Hard coded name.
            return cst.RemoveFromParent()  # Here we have a bug!
        return updated_node


def refactor_string(source: str) -> str:
    transformer = RemoveUnusedImportTransformer()
    cst_tree = cst.parse_module(source)
    return cst_tree.visit(transformer).code


if __name__ == "__main__":
    refactored_code = refactor_string(
        """# Python imports
import os

from typing import Union

# Local imports
from .pgen2 import token
from .pgen2 import driver

from .pgen2.grammar import Grammar

os, Union, driver, Grammar"""
    )
    print(refactored_code)

# Python imports
import os

from typing import Union
from .pgen2 import driver

from .pgen2.grammar import Grammar

os, Union, driver, Grammar