add format string support (#365)

Format strings require support for two new nodes, FormattedValue, respectively JoinedStr.

6 files changed, 61 insertions, 5 deletions

diff --git a/astroid/as_string.py b/astroid/as_string.py
index 3d669c42..8eaebc15 100644
--- a/astroid/as_string.py
+++ b/astroid/as_string.py
@@ -481,6 +481,19 @@ class AsStringVisitor3(AsStringVisitor):
     def visit_asyncfor(self, node):
         return 'async %s' % self.visit_for(node)
 
+    def visit_joinedstr(self, node):
+        # Special treatment for constants,
+        # as we want to join literals not reprs
+        string = ''.join(
+            value.value if type(value).__name__ == 'Const'
+            else value.accept(self)
+            for value in node.values
+        )
+        return "f'%s'" % string
+
+    def visit_formattedvalue(self, node):
+        return '{%s}' % node.value.accept(self)
+
 
 def _import_string(names):
     """return a list of (name, asname) formatted as a string"""
@@ -490,7 +503,7 @@ def _import_string(names):
             _names.append('%s as %s' % (name, asname))
         else:
             _names.append(name)
-    return  ', '.join(_names)
+    return ', '.join(_names)
 
 
 if sys.version_info >= (3, 0):
diff --git a/astroid/node_classes.py b/astroid/node_classes.py
index 42fdf7cc..98731583 100644
--- a/astroid/node_classes.py
+++ b/astroid/node_classes.py
@@ -1863,6 +1863,28 @@ class DictUnpack(NodeNG):
     """Represents the unpacking of dicts into dicts using PEP 448."""
 
 
+class FormattedValue(NodeNG):
+    """Represents a PEP 498 format string."""
+    _astroid_fields = ('value', 'format_spec')
+    value = None
+    conversion = None
+    format_spec = None
+
+    def postinit(self, value, conversion=None, format_spec=None):
+        self.value = value
+        self.conversion = conversion
+        self.format_spec = format_spec
+
+
+class JoinedStr(NodeNG):
+    """Represents a list of string expressions to be joined."""
+    _astroid_fields = ('values',)
+    value = None
+
+    def postinit(self, values=None):
+        self.values = values
+
+
 class Unknown(NodeNG):
     '''This node represents a node in a constructed AST where
     introspection is not possible.  At the moment, it's only used in
diff --git a/astroid/nodes.py b/astroid/nodes.py
index 1c279ccf..3397294b 100644
--- a/astroid/nodes.py
+++ b/astroid/nodes.py
@@ -37,6 +37,7 @@ from astroid.node_classes import (
     TryExcept, TryFinally, Tuple, UnaryOp, While, With, Yield, YieldFrom,
     const_factory,
     AsyncFor, Await, AsyncWith,
+    FormattedValue, JoinedStr,
     # Backwards-compatibility aliases
     Backquote, Discard, AssName, AssAttr, Getattr, CallFunc, From,
     # Node not present in the builtin ast module.
@@ -75,4 +76,5 @@ ALL_NODE_CLASSES = (
     UnaryOp,
     While, With,
     Yield, YieldFrom,
+    FormattedValue, JoinedStr,
     )
diff --git a/astroid/rebuilder.py b/astroid/rebuilder.py
index d80033fe..91774cfd 100644
--- a/astroid/rebuilder.py
+++ b/astroid/rebuilder.py
@@ -869,6 +869,18 @@ class TreeRebuilder3(TreeRebuilder):
     def visit_asyncwith(self, node, parent):
         return self._visit_with(nodes.AsyncWith, node, parent)
 
+    def visit_joinedstr(self, node, parent):
+        newnode = nodes.JoinedStr(node.lineno, node.col_offset, parent)
+        newnode.postinit([self.visit(child, newnode)
+                          for child in node.values])
+        return newnode
+
+    def visit_formattedvalue(self, node, parent):
+        newnode = nodes.FormattedValue(node.lineno, node.col_offset, parent)
+        newnode.postinit(self.visit(node.value, newnode),
+                         node.conversion,
+                         _visit_or_none(node, 'format_spec', self, newnode))
+        return newnode
 
 if sys.version_info >= (3, 0):
     TreeRebuilder = TreeRebuilder3
diff --git a/astroid/tests/unittest_python3.py b/astroid/tests/unittest_python3.py
index ad8e57ba..e555bb48 100644
--- a/astroid/tests/unittest_python3.py
+++ b/astroid/tests/unittest_python3.py
@@ -87,7 +87,7 @@ class Python3TC(unittest.TestCase):
     @require_version('3.0')
     def test_metaclass_imported(self):
         astroid = self.builder.string_build(dedent("""
-        from abc import ABCMeta 
+        from abc import ABCMeta
         class Test(metaclass=ABCMeta): pass"""))
         klass = astroid.body[1]
 
@@ -98,7 +98,7 @@ class Python3TC(unittest.TestCase):
     @require_version('3.0')
     def test_as_string(self):
         body = dedent("""
-        from abc import ABCMeta 
+        from abc import ABCMeta
         class Test(metaclass=ABCMeta): pass""")
         astroid = self.builder.string_build(body)
         klass = astroid.body[1]
@@ -239,6 +239,13 @@ class Python3TC(unittest.TestCase):
             self.assertIsInstance(value, nodes.Const)
             self.assertEqual(value.value, expected)
 
+    @require_version('3.6')
+    def test_format_string(self):
+        code = "f'{greetings} {person}'"
+        node = extract_node(code)
+        self.assertEqual(node.as_string(), code)
+
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/tox.ini b/tox.ini
index db8bf056..982eb388 100644
--- a/tox.ini
+++ b/tox.ini
@@ -1,5 +1,5 @@
 [tox]
-envlist = py27, py33, py34, py35, pypy, jython, pylint
+envlist = py27, py33, py34, py35, py36, pypy, jython, pylint
 skip_missing_interpreters = true
 
 [testenv:pylint]
@@ -10,7 +10,7 @@ deps =
   py27,py33,pypy,jython: enum34
   lazy-object-proxy
   nose
-  py27,py33,py34,py35: numpy
+  py27,py33,py34,py35,py36: numpy
   pytest
   python-dateutil
   py27,py33,pypy,jython: singledispatch