brain_io.py 1.5 KB

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