diff --git a/fire/parser.py b/fire/parser.py index a335cc2c..a8b5c6a8 100644 --- a/fire/parser.py +++ b/fire/parser.py @@ -96,7 +96,7 @@ def _LiteralEval(value): SyntaxError: If the value string has a syntax error. """ root = ast.parse(value, mode='eval') - if isinstance(root.body, ast.BinOp): + if _ContainsDisallowedBinOp(root): raise ValueError(value) for node in ast.walk(root): @@ -116,6 +116,30 @@ def _LiteralEval(value): return ast.literal_eval(root) +def _ContainsDisallowedBinOp(node): + """Returns whether the AST contains BinOps other than complex literals.""" + for child in ast.walk(node): + if isinstance(child, ast.BinOp) and not _IsComplexLiteralBinOp(child): + return True + return False + + +def _IsComplexLiteralBinOp(node): + """Returns whether a BinOp is used to express a complex literal, e.g. 2+3j.""" + if not isinstance(node.op, (ast.Add, ast.Sub)): + return False + return (_HasImaginaryLiteral(node.left) or + _HasImaginaryLiteral(node.right)) + + +def _HasImaginaryLiteral(node): + if isinstance(node, ast.Constant) and isinstance(node.value, complex): + return True + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub): + return _HasImaginaryLiteral(node.operand) + return False + + 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..9b1041e0 100644 --- a/fire/parser_test.py +++ b/fire/parser_test.py @@ -135,6 +135,12 @@ def testDefaultParseValueSyntaxError(self): def testDefaultParseValueIgnoreBinOp(self): self.assertEqual(parser.DefaultParseValue('2017-10-10'), '2017-10-10') self.assertEqual(parser.DefaultParseValue('1+1'), '1+1') + self.assertEqual(parser.DefaultParseValue('[1+1]'), '[1+1]') + self.assertEqual(parser.DefaultParseValue('{key: 1+1}'), '{key: 1+1}') + + def testDefaultParseValueComplexLiterals(self): + self.assertEqual(parser.DefaultParseValue('2+3j'), 2 + 3j) + self.assertEqual(parser.DefaultParseValue('(2+3j)'), 2 + 3j) if __name__ == '__main__': testutils.main()