Useful decorator

[James Stroud]

Paul Rubin wrote:
> In the course of writing a bunch of generator-oriented code I kept
> wanting to temporarily truncate the output of generators for debugging,
> i.e. I might have a function like
> 
>     def generate_bazillion_items():
>        for line in bazillion_line_file:
>           yield itemify(line)
> 
> where I wanted to test the program on just the first 20 lines of the
> file.  After a number of rounds of editing the program to wrap an
> xrange(20) loop I went to using islice(bazillion_line_file, 20), but
> even more handy was to write this decorator:
> 
>     def truncate(n): 
>          from itertools import islice
>          def trunc(gen):  # truncate generator to n items
>             return lambda *a,**kw: islice(gen(*a,**kw), n)
>          return trunc
> 
> Then to chop bazillion_items to 20 items, I just write:
> 
>     @truncate(20)
>     def generate_bazillion_items():
>        for line in bazillion_line_file:
>           yield itemify(line)
> 
> When I want to run the whole file, I comment out the @truncate line,
> and if I want to debug again, I just uncomment it.

Thes may be overkill, but how about allowing None as a parameter to 
truncate, to allow for programmatic truncation?

     # top of module
     if DEBUG:
       FILETRUNC = 20
     else:
       FILETRUNC = None

     # imported, etc.
     def truncate(n):
          if n is None:
            return lambda gen: gen
          from itertools import islice
          def trunc(gen):  # truncate generator to n items
             return lambda *a,**kw: islice(gen(*a,**kw), n)
          return trunc

     # in the body somewhere
     @truncate(FILETRUNC)
     def generate_bazillion_items():
       # etc.

James