Robert Kern avatar Robert Kern committed bc8e1e6

ENH: Support PEP-342 generators in Python 2.5+, too.

Comments (0)

Files changed (2)

         from profile import Profile
 
 
+CO_GENERATOR = 0x0020
+def is_generator(f):
+    """ Return True if a function is a generator.
+    """
+    isgen = (f.func_code.co_flags & CO_GENERATOR) != 0 
+    return isgen
+
+# FIXME: refactor this stuff so that both LineProfiler and ContextualProfile can
+# use the same implementation.
+# Code to exec inside of ContextualProfile.__call__ to support PEP-342-style
+# generators in Python 2.5+.
+pep342_gen_wrapper = '''
+def wrap_generator(self, func):
+    """ Wrap a generator to profile it.
+    """
+    def f(*args, **kwds):
+        g = func(*args, **kwds)
+        # The first iterate will not be a .send()
+        self.enable_by_count()
+        try:
+            item = g.next()
+        finally:
+            self.disable_by_count()
+        input = (yield item)
+        # But any following one might be.
+        while True:
+            self.enable_by_count()
+            try:
+                item = g.send(input)
+            finally:
+                self.disable_by_count()
+            input = (yield item)
+    return f
+'''
+
 class ContextualProfile(Profile):
     """ A subclass of Profile that adds a context manager for Python
     2.5 with: statements and a decorator.
         """ Decorate a function to start the profiler on function entry and stop
         it on function exit.
         """
+        # FIXME: refactor this into a utility function so that both it and
+        # line_profiler can use it.
+        self.add_function(func)
+        if is_generator(func):
+            f = self.wrap_generator()
+        else:
+            f = self.wrap_function()
+        f.__module__ = func.__module__
+        f.__name__ = func.__name__
+        f.__doc__ = func.__doc__
+        f.__dict__.update(getattr(func, '__dict__', {}))
+        return f
+
+    if sys.version_info[:2] >= (2,5):
+        # Delay compilation because the syntax is not compatible with older
+        # Python versions.
+        exec pep342_gen_wrapper
+    else:
+        def wrap_generator(self, func):
+            """ Wrap a generator to profile it.
+            """
+            def f(*args, **kwds):
+                g = func(*args, **kwds)
+                while True:
+                    self.enable_by_count()
+                    try:
+                        item = g.next()
+                    finally:
+                        self.disable_by_count()
+                    yield item
+            return f
+
+    def wrap_function(self, func):
+        """ Wrap a function to profile it.
+        """
         def f(*args, **kwds):
             self.enable_by_count()
             try:
             finally:
                 self.disable_by_count()
             return result
-        f.__name__ = func.__name__
-        f.__doc__ = func.__doc__
-        f.__dict__.update(func.__dict__)
         return f
 
     def __enter__(self):
     isgen = (f.func_code.co_flags & CO_GENERATOR) != 0 
     return isgen
 
+# Code to exec inside of LineProfiler.__call__ to support PEP-342-style
+# generators in Python 2.5+.
+pep342_gen_wrapper = '''
+def wrap_generator(self, func):
+    """ Wrap a generator to profile it.
+    """
+    def f(*args, **kwds):
+        g = func(*args, **kwds)
+        # The first iterate will not be a .send()
+        self.enable_by_count()
+        try:
+            item = g.next()
+        finally:
+            self.disable_by_count()
+        input = (yield item)
+        # But any following one might be.
+        while True:
+            self.enable_by_count()
+            try:
+                item = g.send(input)
+            finally:
+                self.disable_by_count()
+            input = (yield item)
+    return f
+'''
+
 class LineProfiler(CLineProfiler):
     """ A profiler that records the execution times of individual lines.
     """
         """
         self.add_function(func)
         if is_generator(func):
-            def f(*args, **kwds):
-                self.enable_by_count()
-                try:
-                    g = func(*args, **kwds)
-                finally:
-                    self.disable_by_count()
-                while True:
-                    self.enable_by_count()
-                    try:
-                        yield g.next()
-                    finally:
-                        self.disable_by_count()
+            f = self.wrap_generator(func)
         else:
-            # Just a regular function.
-            def f(*args, **kwds):
-                self.enable_by_count()
-                try:
-                    result = func(*args, **kwds)
-                finally:
-                    self.disable_by_count()
-                return result
+            f = self.wrap_function(func)
         f.__module__ = func.__module__
         f.__name__ = func.__name__
         f.__doc__ = func.__doc__
         f.__dict__.update(getattr(func, '__dict__', {}))
         return f
 
+    if sys.version_info[:2] >= (2,5):
+        # Delay compilation because the syntax is not compatible with older
+        # Python versions.
+        exec pep342_gen_wrapper
+    else:
+        def wrap_generator(self, func):
+            """ Wrap a generator to profile it.
+            """
+            def f(*args, **kwds):
+                g = func(*args, **kwds)
+                while True:
+                    self.enable_by_count()
+                    try:
+                        item = g.next()
+                    finally:
+                        self.disable_by_count()
+                    yield item
+            return f
+
+    def wrap_function(self, func):
+        """ Wrap a function to profile it.
+        """
+        def f(*args, **kwds):
+            self.enable_by_count()
+            try:
+                result = func(*args, **kwds)
+            finally:
+                self.disable_by_count()
+            return result
+        return f
+
     def dump_stats(self, filename):
         """ Dump a representation of the data to a file as a pickled LineStats
         object from `get_stats()`.
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.