Merge pull request #142 from bcaller/bc-f-string

Add f-strings to label visitor

pyt/helper_visitors/label_visitor.py

@@ -288,3 +288,35 @@ def visit_Name(self, node):
 
     def visit_Str(self, node):
         self.result += "'" + node.s + "'"
+
+    def visit_joined_str(self, node, surround=True):
+        for val in node.values:
+            if isinstance(val, ast.Str):
+                self.result += val.s
+            else:
+                self.visit(val)
+
+    def visit_JoinedStr(self, node):
+        """
+            JoinedStr(expr* values)
+        """
+        self.result += "f\'"
+        self.visit_joined_str(node)
+        self.result += "'"
+
+    def visit_FormattedValue(self, node):
+        """
+            FormattedValue(expr value, int? conversion, expr? format_spec)
+        """
+        self.result += '{'
+        self.visit(node.value)
+        self.result += {
+            -1: '',     # no formatting
+            97: '!a',   # ascii formatting
+            114: '!r',  # repr formatting
+            115: '!s',  # string formatting
+        }[node.conversion]
+        if node.format_spec:
+            self.result += ':'
+            self.visit_joined_str(node.format_spec)
+        self.result += '}'

tests/helper_visitors/label_visitor_test.py

@@ -71,3 +71,11 @@ def test_list_one_element(self):
     def test_list_two_elements(self):
         label = self.perform_labeling_on_expression('[1, 2]')
         self.assertEqual(label.result, '[1, 2]')
+
+    def test_joined_str(self):
+        label = self.perform_labeling_on_expression('f"a{f(b)}{c}d"')
+        self.assertEqual(label.result, 'f\'a{f(b)}{c}d\'')
+
+    def test_joined_str_with_format_spec(self):
+        label = self.perform_labeling_on_expression('f"a{b!s:.{length}}"')
+        self.assertEqual(label.result, 'f\'a{b!s:.{length}}\'')