Proper way to wrap and unwrap a Python function?

Question

I'm writing wrappers for the Python print function, but my question is more general - having wrapped a function, what's the proper way to un-wrap it?

This works, but I have two concerns about it:

class Iprint():

    def __init__(self, tab=4, level=0):
        ''' Indented printer class.
                tab controls number of spaces per indentation level (equiv. to tabstops)
                level is the indentation level (0=none)'''

        global print

        self.tab = tab
        self.level = level

        self.old_print = print
        print = self.print

    def print(self, *args, end="\n", **kwargs):

        indent = self.tab * self.level

        self.old_print(" "*indent, end="", **kwargs)
        self.old_print(*args, end=end, **kwargs)

indent = Iprint()
indent.level = 3

print("this should be indented")
print = indent.old_print
print("this shouldn't be indented")

My two concerns:

1. What happens if there's a second instantiation of the Iprint() class? This seems awkward and maybe something I ought to prevent - but how?

2. The 2nd to last line print = indent.old_print "unwraps" the print function, returning it to it's original function. This seems awkward too - what if it's forgotten?

I could do it in an __exit__ method but that would restrict the use of this to a with block - I think. Is there a better way?

What's the Pythonic way to do this?

(I also should mention that I anticipate having nested wrappers, which I thinks makes doing this properly more important...)

Answer 1

What it seems you are really trying to do here is find a way to override the builtin print function in a "pythonic" way.

While there is a way to do this, I do have a word of caution. One of the rules of "pythonic code" is

Explicit is better than implicit.

Overwriting print is inherently an implicit solution, and it would be more "pythonic" to allow for a custom print function to solve your needs.

However, let's assume we are talking about a use case where the best option available is to override print. For example, lets say you want to indent the output from the help() function.

You could override print directly, but you run the risk of causing unexpected changes you can't see.

For example:

def function_that_prints():
    log_file = open("log_file.txt", "a")
    print("This should be indented")
    print("internally logging something", file = log_file)
    log_file.close()
    

indent = Iprint()
indent.level = 3
function_that_prints() # now this internal log_file.txt has been corrupted
print = indent.old_print

This is bad, since presumably you just meant to change the output that is printed on screen, and not internal places where print may or may not be used. Instead, you should just override the stdout, not print.

Python now includes a utility to do this called contextlib.redirect_stdout() documented here.

An implementation may look like this:

import io
import sys
import contextlib

class StreamIndenter(io.TextIOBase):
    # io.TextIOBase provides some base functions, such as writelines()

    def __init__(self, tab = 4, level = 1, newline = "\n", stream = sys.stdout):
        """Printer that adds an indent at the start of each line"""
        self.tab = tab
        self.level = level
        self.stream = stream
        self.newline = newline
        self.linestart = True

    def write(self, buf, *args, **kwargs):
        if self.closed:
            raise ValueError("write to closed file")

        if not buf:
            # Quietly ignore printing nothing
            # prevents an issue with print(end='')
            return 

        indent = " " * (self.tab * self.level)

        if self.linestart:
            # The previous line has ended. Indent this one
            self.stream.write(indent)

        # Memorize if this ends with a newline
        if buf.endswith(self.newline):
            self.linestart = True

            # Don't replace the last newline, as the indent would double
            buf = buf[:-len(self.newline)]
            self.stream.write(buf.replace(self.newline, self.newline + indent))
            self.stream.write(self.newline)
        else:
            # Does not end on a newline
            self.linestart = False
            self.stream.write(buf.replace(self.newline, self.newline + indent))

    # Pass some calls to internal stream
    @property
    def writable(self):
        return self.stream.writable

    @property
    def encoding(self):
        return self.stream.encoding

    @property
    def name(self):
        return self.stream.name


with contextlib.redirect_stdout(StreamIndenter()) as indent:
    indent.level = 2
    print("this should be indented")
print("this shouldn't be indented")

Overriding print this way both doesn't corrupt other uses of print and allows for proper handling of more complicated usages.

For example:

with contextlib.redirect_stdout(StreamIndenter()) as indent:
    indent.level = 2
    print("this should be indented")

    indent.level = 3
    print("more indented")

    indent.level = 2
    for c in "hello world\n": print(c, end='')
    print()
    print("\n", end='')
    print(end = '')
    
print("this shouldn't be indented")

Formats correctly as:

        this should be indented
            more indented
        hello world
        
        
this shouldn't be indented

Answer 2

I think I've solved this - at least to my own satisfaction. Here I've called the class T (for test):

class T():

    old_print = None

    def __init__(self, tab=4, level=0):
        ''' Indented printer class.
                tab controls number of spaces per indentation level (equiv. to tabstops)
                level is the indentation level (0=none)'''

        T.tab = tab
        T.level = level

        self.__enter__()


    def print(self, *args, end="\n", **kwargs):

        indent = T.tab * T.level

        T.old_print(" "*indent, end="", **kwargs)
        T.old_print(*args, end=end, **kwargs)


    def close(self):

        if T.old_print is not None:

            global print
            print = T.old_print
            T.old_print = None

    def __enter__(self):
        if T.old_print is None:

            global print
            T.old_print = print
            print = self.print

    def __exit__(self, exception_type, exception_value, exception_traceback):
        self.close()


print("this should NOT be indented")

i = T(level=1)

print("level 1")

i2 = T(level=2)

print("level 2")

i.close()

print("this should not be indented")

i3 = T(level=3)

print("level 3")

i2.close()

print("not indented")

with i:
    print("i")

print("after i")

with T(level=3):
    print("T(level=3)")

print("after T(level=3)")

It silently forces a single (functional) instance of the class, regardless of how many times T() is called, as @MichaelButscher suggested (thanks; that was the most helpful comment by far).

It works cleanly with WITH blocks, and you can manually call the close method if not using WITH blocks.

The output is, as expected:

this should NOT be indented
    level 1
        level 2
this should not be indented
            level 3
not indented
            i
after i
            T(level=3)
after T(level=3)