Writing Custom Modules

Let’s create a python file mymodule.py with the following content.

%%file mymodule.py

print("BEGIN mymodule")
x = 2

def add(a, b):
    return a + b

print(add(3, 4))
print("END mymodule")

Overwriting mymodule.py

Let’s see what happens when we run this file as a script:

!python mymodule.py

BEGIN mymodule
7
END mymodule

Let’s see what happens when we import this file as a module.

import mymodule

BEGIN mymodule
7
END mymodule

mymodule.x

mymodule.add(10, 20)

As you can see a python file can be imported as a module and the code in that file gets executed, just the same way when it is run as as script.

The output would be the same if we import the module from another python file or the program is passed directly from python command-line ¹.

!python -c "import mymodule; print(mymodule.add(10, 20))"

BEGIN mymodule
7
END mymodule
30

Reimporting a module

The code of the module is executed only once even if the module is imported more than once.

import mymodule
import mymodule

That means, if you make some changes the file of a module after it is imported, you’ll not see the changes even if you import it again.

This is especially annoying when you are working in a Jupyter notebook.

The work-around to that issue is to reimport a module using importlib.

import mymodule

# and make some changes to mymodule.py

import importlib
importlib.reload(mymodule)

BEGIN mymodule
7
END mymodule

<module 'mymodule' from '/home/anand/trainings/2023/perfios-python/book/reference/mymodule.py'>

The `name` magic variable

Python has a special variable __name__, which is set by the runtime.

The value of the __name__ will be different in each module or Python file. If a Python file is executed as a script, the value of __name__ will be "__main__". When a python file is imported as a module, the value of __name__ would be the module name.

Let’s try this with an example.

%%file mymodule2.py

x = 2

def add(a, b):
    return a+b

print(__name__)
print(add(3, 4))

Overwriting mymodule2.py

!python mymodule2.py

__main__
7

The value of __name__ is "__main__" becasuse the file is executed as a script.

!python -c "import mymodule2"

mymodule2
7

Since we have now imported the file as a mofule, the value of __name__ is "mymodule2", the name of the module.

If you notice, the module is printing the value of add(3, 4) even when the file is imported as a module, which is not desirable. How do we stop it?

%%file mymodule3.py

x = 2

def add(a, b):
    return a+b

# run the following only when this file is run as a script
if __name__ == "__main__":
    print(add(3, 4))

Overwriting mymodule3.py

We could check if the program is run as a script by using the condition if __name__ == "__main__": and then only run whatever we want to run it when the file is used as a script.

Example: square module

%%file sq.py
import sys

def square(x):
    return x*x

def main():
    n = int(sys.argv[1])
    print(square(n))

if __name__ == "__main__":
    main()

Overwriting sq.py

!python sq.py 5

import sq
sq.square(5)

Docstrings

Let’s see how to find help with Python modules and functions and how to make our modules provide such help.

import os

help(os.listdir)

Help on built-in function listdir in module posix:

listdir(path=None)
    Return a list containing the names of the files in the directory.
    
    path can be specified as either str, bytes, or a path-like object.  If path is bytes,
      the filenames returned will also be bytes; in all other circumstances
      the filenames returned will be str.
    If path is None, uses the path='.'.
    On some platforms, path may also be specified as an open file descriptor;\
      the file descriptor must refer to a directory.
      If this functionality is unavailable, using it raises NotImplementedError.
    
    The list is in arbitrary order.  It does not include the special
    entries '.' and '..' even if they are present in the directory.

Adding docstrings to a function

We can add docstrings to a function by adding a string as the first line of the function.

def add(x, y):
    "Adds two numbers"
    return x+y

add(3, 4)

help(add)

Help on function add in module __main__:

add(x, y)
    Adds two numbers

Typically, triple quotes are used to write docstrings as they allow writing docstring in multiple lines. Also, it is often useful to include an example of how the function can be used.

def add(x, y):
    """Adds two numbers.

        >>> add(3, 4)
        7
    """
    return x+y

help(add)

Help on function add in module __main__:

add(x, y)
    Adds two numbers.
    
    >>> add(3, 4)
    7

Using typehints

Python3 supports adding typehints to funtions. It is often very useful to include typehints for function as they are shown as part of the documentation.

def add(x: int, y: int) -> int:
    """Adds two numbers.

        >>> add(3, 4)
        7
    """
    return x+y

help(add)

Help on function add in module __main__:

add(x: int, y: int) -> int
    Adds two numbers.
    
    >>> add(3, 4)
    7

Adding docstrings to a module

Docstrings can also be added a module, just like how they are added to a function. Just add a string at the beginning of the file.

Let’s add docstrings to the sq module that we wrote earlier.

%%file sq.py
"""
The square module.

The square module provides a function to compute square of a number.

This can also be used as a script.

USAGE:

    $ python sq.py 4
    16
"""
import sys

def square(n: int) -> int:
    """Computes square of a number.
    
        >>> square(4)
        16
    """
    return n*n

def main():
    n = int(sys.argv[1])
    print(square(n))
    
if __name__ == "__main__":
    main()

Overwriting sq.py

We may have to reimport the module as it may have been imported earlier.

import sq
import importlib

importlib.reload(sq)

<module 'sq' from '/home/anand/trainings/2023/perfios-python/book/reference/sq.py'>

The reload was required because we were modifying the file of a module that was already imported earlier. You don’t have to do this, if you make a new module.

Now that we have the docstrings ready for the sq module, let’s look at the help.

help(sq)

Help on module sq:

NAME
    sq - The square module.

DESCRIPTION
    The square module provides a function to compute square of a number.
    
    This can also be used as a script.
    
    USAGE:
    
        $ python sq.py 4
        16

FUNCTIONS
    main()
    
    square(n: int) -> int
        Computes square of a number.
        
        >>> square(4)
        16

FILE
    /home/anand/trainings/2023/perfios-python/book/reference/sq.py

Or help just on one function.

help(sq.square)

Help on function square in module sq:

square(n: int) -> int
    Computes square of a number.
    
    >>> square(4)
    16

Footnotes

The python command-line flag -c allows the program directly from the command line. It is handy to quickly try some small snippets of code. Then the code is multiple statements, they are usually seperated by a semi-comon.↩︎