[pypy-commit] pypy callfamily: replace bk.enter()/bk.leave() pairs with a context manager

[rlamy]

Author: Ronan Lamy <ronan.lamy at gmail.com>
Branch: callfamily
Changeset: r76798:1ebe0debf9ff
Date: 2015-04-15 18:48 +0100
http://bitbucket.org/pypy/pypy/changeset/1ebe0debf9ff/

Log:	replace bk.enter()/bk.leave() pairs with a context manager

diff --git a/rpython/annotator/annrpython.py b/rpython/annotator/annrpython.py
--- a/rpython/annotator/annrpython.py
+++ b/rpython/annotator/annrpython.py
@@ -407,8 +407,7 @@
             i = 0
             while i < len(block.operations):
                 op = block.operations[i]
-                self.bookkeeper.enter((graph, block, i))
-                try:
+                with self.bookkeeper.at_position((graph, block, i)):
                     new_ops = op.transform(self)
                     if new_ops is not None:
                         block.operations[i:i+1] = new_ops
@@ -417,8 +416,6 @@
                         new_ops[-1].result = op.result
                         op = new_ops[0]
                     self.consider_op(op)
-                finally:
-                    self.bookkeeper.leave()
                 i += 1
 
         except BlockedInference as e:
diff --git a/rpython/annotator/bookkeeper.py b/rpython/annotator/bookkeeper.py
--- a/rpython/annotator/bookkeeper.py
+++ b/rpython/annotator/bookkeeper.py
@@ -5,6 +5,7 @@
 from __future__ import absolute_import
 
 import sys, types, inspect, weakref
+from contextlib import contextmanager
 
 from rpython.flowspace.model import Constant
 from rpython.annotator.model import (SomeOrderedDict,
@@ -87,11 +88,22 @@
         del TLS.bookkeeper
         del self.position_key
 
+    @contextmanager
+    def at_position(self, pos):
+        """A context manager calling `self.enter()` and `self.leave()`"""
+        if hasattr(self, 'position_key') and pos is None:
+            yield
+            return
+        self.enter(pos)
+        try:
+            yield
+        finally:
+            self.leave()
+
     def compute_at_fixpoint(self):
         # getbookkeeper() needs to work during this function, so provide
         # one with a dummy position
-        self.enter(None)
-        try:
+        with self.at_position(None):
             for call_op in self.annotator.call_sites():
                 self.consider_call_site(call_op)
 
@@ -99,8 +111,6 @@
                 args = simple_args(args_s)
                 pbc.consider_call_site(args, s_ImpossibleValue, None)
             self.emulated_pbc_calls = {}
-        finally:
-            self.leave()
 
     def check_no_flags_on_instances(self):
         # sanity check: no flags attached to heap stored instances
@@ -480,7 +490,6 @@
         """Analyse a call to a SomePBC() with the given args (list of
         annotations).
         """
-        descs = list(pbc.descriptions)
         if emulated is None:
             whence = self.position_key
             # fish the existing annotation for the result variable,
@@ -505,10 +514,7 @@
         return s_result
 
     def emulate_pbc_call(self, unique_key, pbc, args_s, replace=[], callback=None):
-        emulate_enter = not hasattr(self, 'position_key')
-        if emulate_enter:
-            self.enter(None)
-        try:
+        with self.at_position(None):
             emulated_pbc_calls = self.emulated_pbc_calls
             prev = [unique_key]
             prev.extend(replace)
@@ -523,9 +529,6 @@
             else:
                 emulated = callback
             return self.pbc_call(pbc, args, emulated=emulated)
-        finally:
-            if emulate_enter:
-                self.leave()
 
     def _find_current_op(self, opname=None, arity=None, pos=None, s_type=None):
         """ Find operation that is currently being annotated. Do some
diff --git a/rpython/annotator/description.py b/rpython/annotator/description.py
--- a/rpython/annotator/description.py
+++ b/rpython/annotator/description.py
@@ -361,11 +361,8 @@
     @staticmethod
     def variant_for_call_site(bookkeeper, family, descs, args, op):
         shape = rawshape(args)
-        bookkeeper.enter(None)
-        try:
+        with bookkeeper.at_position(None):
             row = FunctionDesc.row_to_consider(descs, args, op)
-        finally:
-            bookkeeper.leave()
         index = family.calltable_lookup_row(shape, row)
         return shape, index