brain_io.py 1.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445
  1. # Copyright (c) 2016, 2018, 2020 Claudiu Popa <pcmanticore@gmail.com>
  2. # Copyright (c) 2020-2021 hippo91 <guillaume.peillex@gmail.com>
  3. # Copyright (c) 2021 Pierre Sassoulas <pierre.sassoulas@gmail.com>
  4. # Copyright (c) 2021 Marc Mueller <30130371+cdce8p@users.noreply.github.com>
  5. # Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html
  6. # For details: https://github.com/PyCQA/astroid/blob/main/LICENSE
  7. """Astroid brain hints for some of the _io C objects."""
  8. from astroid.manager import AstroidManager
  9. from astroid.nodes import ClassDef
  10. BUFFERED = {"BufferedWriter", "BufferedReader"}
  11. TextIOWrapper = "TextIOWrapper"
  12. FileIO = "FileIO"
  13. BufferedWriter = "BufferedWriter"
  14. def _generic_io_transform(node, name, cls):
  15. """Transform the given name, by adding the given *class* as a member of the node."""
  16. io_module = AstroidManager().ast_from_module_name("_io")
  17. attribute_object = io_module[cls]
  18. instance = attribute_object.instantiate_class()
  19. node.locals[name] = [instance]
  20. def _transform_text_io_wrapper(node):
  21. # This is not always correct, since it can vary with the type of the descriptor,
  22. # being stdout, stderr or stdin. But we cannot get access to the name of the
  23. # stream, which is why we are using the BufferedWriter class as a default
  24. # value
  25. return _generic_io_transform(node, name="buffer", cls=BufferedWriter)
  26. def _transform_buffered(node):
  27. return _generic_io_transform(node, name="raw", cls=FileIO)
  28. AstroidManager().register_transform(
  29. ClassDef, _transform_buffered, lambda node: node.name in BUFFERED
  30. )
  31. AstroidManager().register_transform(
  32. ClassDef, _transform_text_io_wrapper, lambda node: node.name == TextIOWrapper
  33. )