How to make your script better

We have all been there. The first time you realize you know enough programming language that you want to make something.  One of the best things the Python community has is that it encourages the newcomers to not only build amazing things, but to share them with everyone, and thus the community gets better.

However, I have been noticing a pattern in some of these projects. You know which ones I am talking about. Those projects with 2 or 3 different python files, where the README provides a few hints on how to use it and has phrases like I will add x and y functionalities, and the main entrypoint has the same name as the repository.

The thing is, some of this projects are awesome, and they have some kind of logic that I would like to use without having to think about it for myself.

Here are 3 different ways to set up these scripts, each one being (subjectively of course) better than the previous one.

Let's say we want to build a super awesome script, this script will have one functionality. It will get user input as an integer. perform a simple calculation with that integer, then print the result.

Implementation 1

#!/usr/bin/env python

"""
Super awesome script  
Asks the user for a number:  
 - If the number is less or equal to 100, it returns the 1st tetration of the number (power of itself)
 - else, it returns the number squared
"""

__version__ = '0.1'

if __name__ == '__main__':

    while 1:
        user_number = raw_input('Choose a number:\n') #input() in python3
        if user_number.isdigit():
            user_number = int(user_number)
            break
        else:
            print('{} is not a valid number'.format(user_number))

    if user_number > 100:
        print(user_number**2)
    else:
        print(user_number**user_number)

This is the implementation a newcomer to Python would probably write. It works, it asks the user for the input, validates it, performs the calculation and returns the result.

I see two issues here:

  1. There is no separation between the input logic and the calculation logic. Everything is jammed together on the main loop. It works, but it would make it more difficult to modify or expand this code (you would need to read all the code line after line).
  2. We are doing all the input validation by ourselves. Im sure python has some libraries for this.

Of course it does! For the next, more advanced implementation, we will use the awesome builtin module argparse.

Implementation 2

#!/usr/bin/env python

"""
Super awesome script  
Asks the user for a number:  
 - If the number is less or equal to 100, it returns it to the power of itself
 - else, it returns the number squared
"""

import argparse

__version__ = '0.2'


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--number', required=True, type=int,
                        help='number to perform calculation')
    values = parser.parse_args()
    user_number = values.number
    if user_number > 100:
        print(user_number**2)
    else:
        print(user_number**user_number)

On this step we have removed one of the issues we had with the previous implementation. By using argparse we are letting the library handle the validation for ourselves. This script wont run unless the input is valid.

We still have the issue of separation of logic between user input and the calculation. Which brings us to the final implementation...

Implementation 3

#!/usr/bin/env python

"""
Super awesome script  
Asks the user for a number:  
 - If the number is less or equal to 100, it returns it to the power of itself
 - else, it returns the number squared
"""

import argparse

__version__ = '0.3'



def calculation(number):  
    """Performs awesome calculation"""
    if number > 100:
        return number**2
    else:
        return number**number

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--number', required=True, type=int,
                        help='number to perform calculation')
    values = parser.parse_args()
    user_number = values.number
    calculation_result = calculation(user_number)
    print(calculation_result)

On this implementation, we have done two things:

  1. We have put the burden of input validation on a mature library such as argparse
  2. We have separated the input logic from the calculation logic.

This last change has two effects; First, if we realize we want to modify the calculation (lets say we want to change the 100 for 200, we can do that without having to read the whole code. Just need to modify the function calculation and everything else will still work.

And the most important effect. If I read your code and I like the calculation function, now I can import it into my own code!! Think about it, on implementations #1 and #2, the only way to use the script was by doing:

python calculation_script.py --number INTEGER

Now, I have another and much more powerful way of using the code. For example, if I have another script, from that script I can run the calculation logic like:

from calculation_script import calculation

number = 10  
calculation_result = calculation(number)  

Isn't that amazing? The fact that by simply refactoring your main function, now every single other piece of code in the python ecosystem can use your code!

Think about it next, time, it takes very little effort and it makes your script better :)