diff --git a/pyt/helper_visitors/label_visitor.py b/pyt/helper_visitors/label_visitor.py index ab2301b9..f594bff8 100644 --- a/pyt/helper_visitors/label_visitor.py +++ b/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 += '}' diff --git a/tests/helper_visitors/label_visitor_test.py b/tests/helper_visitors/label_visitor_test.py index 0f2d2f7d..e39f5151 100644 --- a/tests/helper_visitors/label_visitor_test.py +++ b/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}}\'')