Robert Kern avatar Robert Kern committed 831dbb1

BUG: correctly profile generators.

Comments (0)

Files changed (1)

 from _line_profiler import LineProfiler as CLineProfiler
 
 
+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
+
 class LineProfiler(CLineProfiler):
     """ A profiler that records the execution times of individual lines.
     """
         """ Decorate a function to start the profiler on function entry and stop
         it on function exit.
         """
-        def f(*args, **kwds):
-            self.add_function(func)
-            self.enable_by_count()
-            try:
-                result = func(*args, **kwds)
-            finally:
-                self.disable_by_count()
-            return result
+        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()
+        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.__module__ = func.__module__
         f.__name__ = func.__name__
         f.__doc__ = func.__doc__
-        f.__dict__.update(func.__dict__)
+        f.__dict__.update(getattr(func, '__dict__', {}))
         return f
 
     def dump_stats(self, filename):
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.