NodeTransformer to replace lazy attach with attach_stub.
| 14 | |
| 15 | |
| 16 | class RewriteAssign(ast.NodeTransformer): |
| 17 | """NodeTransformer to replace lazy attach with attach_stub.""" |
| 18 | |
| 19 | def visit_Assign(self, node): |
| 20 | """Replace lazy attach assignment with stub assignment.""" |
| 21 | if not hasattr(node.targets[0], "dims"): |
| 22 | return node |
| 23 | |
| 24 | ids = [name.id for name in node.targets[0].dims] |
| 25 | if ids == ["__getattr__", "__dir__", "__all__"]: |
| 26 | return ast.parse( |
| 27 | "__getattr__, __dir__, __all__ = lazy.attach_stub(__name__, __file__)\n" |
| 28 | ) |
| 29 | return node |
| 30 | |
| 31 | |
| 32 | pyi_mode = black.Mode(is_pyi=True) |