diff --git a/fire/parser.py b/fire/parser.py index a335cc2c..86761e4b 100644 --- a/fire/parser.py +++ b/fire/parser.py @@ -79,9 +79,10 @@ def DefaultParseValue(value): def _LiteralEval(value): """Parse value as a Python literal, or container of containers and literals. - First the AST of the value is updated so that bare-words are turned into - strings. Then the resulting AST is evaluated as a literal or container of - only containers and literals. + First the AST of the value is checked for binary operations, which are not + literals. Then the AST is updated so that bare-words are turned into strings. + Then the resulting AST is evaluated as a literal or container of only + containers and literals. This allows for the YAML-like syntax {a: b} to represent the dict {'a': 'b'} @@ -96,8 +97,13 @@ def _LiteralEval(value): SyntaxError: If the value string has a syntax error. """ root = ast.parse(value, mode='eval') - if isinstance(root.body, ast.BinOp): - raise ValueError(value) + for node in ast.walk(root): + # Binary operations are not literals, so they are rejected wherever they + # appear. Note that some versions of ast.literal_eval evaluate them + # instead, e.g. turning '[1+1]' into [2]. The only exception is the + or - + # used to write a complex number, which is a literal. + if isinstance(node, ast.BinOp) and not _IsComplexLiteral(node): + raise ValueError(value) for node in ast.walk(root): for field, child in ast.iter_fields(node): @@ -116,6 +122,30 @@ def _LiteralEval(value): return ast.literal_eval(root) +def _IsComplexLiteral(node): + """Returns whether a binary operation expresses a complex number literal. + + Python has no single-token syntax for a complex number with both a real and + an imaginary part, so such numbers are written as the sum or difference of a + real number and an imaginary one, e.g. 1+2j. These are the only binary + operations ast.literal_eval accepts, so they are the only ones Fire accepts. + + Args: + node: A node of type BinOp. + Returns: + True if the node has the form + , with + either + or -, False otherwise. + """ + if not isinstance(node.op, (ast.Add, ast.Sub)): + return False + try: + real = ast.literal_eval(node.left) + imaginary = ast.literal_eval(node.right) + except (TypeError, ValueError): + return False + return isinstance(real, (int, float)) and isinstance(imaginary, complex) + + def _Replacement(node): """Returns a node to use in place of the supplied node in the AST. diff --git a/fire/parser_test.py b/fire/parser_test.py index a404eea2..218ba5cf 100644 --- a/fire/parser_test.py +++ b/fire/parser_test.py @@ -136,5 +136,26 @@ def testDefaultParseValueIgnoreBinOp(self): self.assertEqual(parser.DefaultParseValue('2017-10-10'), '2017-10-10') self.assertEqual(parser.DefaultParseValue('1+1'), '1+1') + def testDefaultParseValueIgnoreNestedBinOp(self): + # Binary operations are not literals, wherever they appear. + self.assertEqual(parser.DefaultParseValue('[1+1]'), '[1+1]') + self.assertEqual(parser.DefaultParseValue('{key: 1+1}'), '{key: 1+1}') + self.assertEqual(parser.DefaultParseValue('(1+1,)'), '(1+1,)') + self.assertEqual(parser.DefaultParseValue('[[2017-10-10]]'), + '[[2017-10-10]]') + self.assertEqual(parser.DefaultParseValue('{1: [a*b]}'), '{1: [a*b]}') + self.assertEqual(parser.DefaultParseValue('1+1+2j'), '1+1+2j') + + def testDefaultParseValueComplexNumbers(self): + # The + and - of a complex number are literals, and are still parsed. + self.assertEqual(parser.DefaultParseValue('2j'), 2j) + self.assertEqual(parser.DefaultParseValue('1+2j'), 1 + 2j) + self.assertEqual(parser.DefaultParseValue('-1.5-2j'), -1.5 - 2j) + self.assertEqual(parser.DefaultParseValue('[1+2j]'), [1 + 2j]) + self.assertEqual(parser.DefaultParseValue('{alpha: 1+2j}'), + {'alpha': 1 + 2j}) + # The real part must come first, as it must for ast.literal_eval. + self.assertEqual(parser.DefaultParseValue('2j+1'), '2j+1') + if __name__ == '__main__': testutils.main()