test more simple_eval functionality

thomasballinger · thomasballinger · commit 6b1d291e023e · 2015-12-14T17:39:19.000-05:00
diff --git a/bpython/autocomplete.py b/bpython/autocomplete.py
@@ -74,6 +74,22 @@
 def after_last_dot(name):
     return name.rstrip('.').rsplit('.')[-1]
 
+def few_enough_underscores(current, match):
+    """Returns whether match should be shown based on current
+
+    if current is _, True if match starts with 0 or 1 underscore
+    if current is __, True regardless of match
+    otherwise True if match does not start with any underscore
+    """
+    if current.startswith('__'):
+        return True
+    elif current.startswith('_') and not match.startswith('__'):
+        return True
+    elif match.startswith('_'):
+        return False
+    else:
+        return True
+
 
 def method_match_simple(word, size, text):
     return word[:size] == text
@@ -255,18 +271,9 @@ def matches(self, cursor_offset, line, **kwargs):
         matches = set(''.join([r.word[:-i], m])
                       for m in self.attr_matches(methodtext, locals_))
 
-        # TODO add open paren for methods via _callable_prefix (or decide not
-        # to) unless the first character is a _ filter out all attributes
-        # starting with a _
-        if r.word.split('.')[-1].startswith('__'):
-            pass
-        elif r.word.split('.')[-1].startswith('_'):
-            matches = set(match for match in matches
-                          if not match.split('.')[-1].startswith('__'))
-        else:
-            matches = set(match for match in matches
-                          if not match.split('.')[-1].startswith('_'))
-        return matches
+        return set(m for m in matches
+                   if few_enough_underscores(r.word.split('.')[-1],
+                                             m.split('.')[-1]))
 
     def locate(self, current_offset, line):
         return lineparts.current_dotted_attribute(current_offset, line)
@@ -470,14 +477,8 @@ def matches(self, cursor_offset, line, **kwargs):
                        # strips leading dot
             matches = [m[1:] for m in self.attr_lookup(obj, '', attr.word)]
 
-        if attr.word.startswith('__'):
-            pass
-        elif attr.word.startswith('_'):
-            matches = set(match for match in matches
-                          if not match.startswith('__'))
-        else:
-            matches = set(match for match in matches
-                          if not match.split('.')[-1].startswith('_'))
+
+        return set(m for m in matches if few_enough_underscores(attr.word, m))
         return matches
 
 
diff --git a/bpython/simpleeval.py b/bpython/simpleeval.py
@@ -168,10 +168,12 @@ def evaluate_current_expression(cursor_offset, line, namespace=None):
     """
     Return evaluated expression to the right of the dot of current attribute.
 
-    build asts from with increasing numbers of characters.
-    Find the biggest valid ast.
-    Once our attribute access is a subtree, stop
+    Only evaluates builtin objects, and do any attribute lookup.
     """
+    # Builds asts from with increasing numbers of characters back from cursor.
+    # Find the biggest valid ast.
+    # Once our attribute access is found, return its .value subtree
+
     if namespace is None:
         namespace = {}
 
@@ -207,6 +209,7 @@ def parse_trees(cursor_offset, line):
 
 
 def evaluate_current_attribute(cursor_offset, line, namespace=None):
+    """Safely evaluates the expression attribute lookup currently occuring on"""
     # this function runs user code in case of custom descriptors,
     # so could fail in any way
 
diff --git a/bpython/test/test_repl.py b/bpython/test/test_repl.py
@@ -230,6 +230,13 @@ def test_issue583(self):
         self.repl.set_docstring()
         self.assertIsNot(self.repl.docstring, None)
 
+    def test_methods_of_expressions(self):
+        self.set_input_line("'a'.capitalize(")
+        self.assertTrue(self.repl.get_args())
+
+        self.set_input_line("(1 + 1).bit_length(")
+        self.assertTrue(self.repl.get_args())
+
 
 class TestGetSource(unittest.TestCase):
     def setUp(self):
diff --git a/bpython/test/test_simpleeval.py b/bpython/test/test_simpleeval.py
@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
 import ast
+import numbers
 
 from bpython.simpleeval import (simple_eval,
                                 evaluate_current_expression,
@@ -60,6 +61,23 @@ def __getattr__(inner_self, attr):
         with self.assertRaises(ValueError):
             simple_eval('a[1]', {'a': SchrodingersDict()})
 
+    def test_operators_on_suspicious_types(self):
+        class Spam(numbers.Number):
+            def __add__(inner_self, other):
+                self.fail("doing attribute lookup might have side effects")
+
+        with self.assertRaises(ValueError):
+            simple_eval('a + 1', {'a': Spam()})
+
+    def test_operators_on_numbers(self):
+        self.assertEqual(simple_eval('-2'), -2)
+        self.assertEqual(simple_eval('1 + 1'), 2)
+        self.assertEqual(simple_eval('a - 2', {'a':1}), -1)
+        with self.assertRaises(ValueError):
+            simple_eval('2 * 3')
+        with self.assertRaises(ValueError):
+            simple_eval('2 ** 3')
+
     def test_function_calls_raise(self):
         with self.assertRaises(ValueError):
             simple_eval('1()')
