Add brain tips for _io.TextIOWrapper's buffer and raw attributes.

3 files changed, 63 insertions, 0 deletions

diff --git a/ChangeLog b/ChangeLog
index 863438e8..f2002e74 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -2,6 +2,8 @@ Change log for the astroid package (used to be astng)
 =====================================================
 
 --
+    * Add brain tips for _io.TextIOWrapper's buffer and raw attributes.
+
     * Add `returns` into the proper order in FunctionDef._astroid_fields
 
       The order is important, since it determines the last child,
diff --git a/astroid/brain/brain_io.py b/astroid/brain/brain_io.py
new file mode 100644
index 00000000..be8d30c8
--- /dev/null
+++ b/astroid/brain/brain_io.py
@@ -0,0 +1,43 @@
+# Copyright (c) 2016 Claudiu Popa <pcmanticore@gmail.com>
+
+# Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html
+# For details: https://github.com/PyCQA/astroid/blob/master/COPYING.LESSER
+
+'''Astroid brain hints for some of the _io C objects.'''
+
+import astroid
+
+
+BUFFERED = {'BufferedWriter', 'BufferedReader'}
+TextIOWrapper = 'TextIOWrapper'
+FileIO = 'FileIO'
+BufferedWriter = 'BufferedWriter'
+
+
+def _generic_io_transform(node, name, cls):
+    '''Transform the given name, by adding the given *class* as a member of the node.'''
+
+    io_module = astroid.MANAGER.ast_from_module_name('_io')
+    attribute_object = io_module[cls]
+    instance = attribute_object.instantiate_class()
+    node.locals[name] = [instance]
+
+
+def _transform_text_io_wrapper(node):
+    # This is not always correct, since it can vary with the type of the descriptor,
+    # being stdout, stderr or stdin. But we cannot get access to the name of the
+    # stream, which is why we are using the BufferedWriter class as a default
+    # value
+    return _generic_io_transform(node, name='buffer', cls=BufferedWriter)
+
+
+def _transform_buffered(node):
+    return _generic_io_transform(node, name='raw', cls=FileIO)
+
+
+astroid.MANAGER.register_transform(astroid.ClassDef,
+                                   _transform_buffered,
+                                   lambda node: node.name in BUFFERED)
+astroid.MANAGER.register_transform(astroid.ClassDef,
+                                   _transform_text_io_wrapper,
+                                   lambda node: node.name == TextIOWrapper)
diff --git a/astroid/tests/unittest_brain.py b/astroid/tests/unittest_brain.py
index a7d27b8c..88d39df3 100644
--- a/astroid/tests/unittest_brain.py
+++ b/astroid/tests/unittest_brain.py
@@ -506,5 +506,23 @@ class PytestBrainTest(unittest.TestCase):
         self.assertIn('mark', module)
 
 
+class IOBrainTest(unittest.TestCase):
+
+    @unittest.skipUnless(six.PY3, 'Needs Python 3 io model')
+    def test_sys_streams(self):
+        for name in {'stdout', 'stderr', 'stdin'}:
+            node = astroid.extract_node('''
+            import sys
+            sys.{}
+            '''.format(name))
+            inferred = next(node.infer())
+            buffer = next(inferred.igetattr('buffer'))
+            self.assertIsInstance(buffer, astroid.Instance)
+            self.assertEqual(buffer.name, 'BufferedWriter')
+            raw = next(buffer.igetattr('raw'))
+            self.assertIsInstance(raw, astroid.Instance)
+            self.assertEqual(raw.name, 'FileIO')
+
+
 if __name__ == '__main__':
     unittest.main()