How to organize code into different parts

When I first learned to code, I wrote all my code in a single file. I didn’t even use functions.

For example, let’s say we had a simple application that allowed the user to enter in some todos and we wrote it to a file when they were done.

My code would look something like this.

import os
import datetime

filename = 'todos1.csv'

todos = []

while True:
    todo = input('What is your todo?')
    if todo == 'done':
        break

    todos.append([todo, int(datetime.datetime.now().timestamp())])

with open(filename, 'a') as f:
    for todo, timestamp in todos:
        f.write(f'{todo},{timestamp}\n')

This works but it’s not great.

There are a number of issues:

  • As we add more functionality, this becomes larger and harder to understand
  • The only way to make sure it works correctly is to manually run it and make sure the file was properly created
  • The high-level functionality is not obvious

How I would improve this today

Instead of having a single block of code, I would break out the different functionality.

The application needs to accept user input, format the input, then write it to a file.

The first thing I would do is isolate the part that accepts user input. Here I’m breaking out a separate function get_user_input which collects user input and returns it in a list.

def get_user_input(prompt):                                                                          
    user_inputs = []                                                                                 
    while True:                                                                                      
        user_input = input(prompt)                                                                   
        if user_input == exit_command:                                                               
            return user_inputs                                                                       
                                                                                                     
        current_timestamp = int(datetime.now().timestamp())                                          
        user_inputs.append((user_input, current_timestamp))  

Next, I would define a function run that controls the main flow of the application.

One tip I try to follow is trying to make the code read like English.

def run(filename, user_inputs, todos_to_csv, write_rows):                                            
    file_contents = todos_to_csv(user_inputs)                                                        
    write_rows(filename, file_contents)                                                              

One thing you may have noticed is that I provided some functions (i.e. todos_to_csv and write_rows) as arguments to run.

The reason I do this is it allows me to isolate the mechanics of run from how those other functions actually work. e.g. if I called write_rows directly, to verify run works I would have to check that a file was actually created.

By providing those functions as arguments, run doesn’t need to care about how those actually work. All it needs to know is that the inputs were transformed and we called the function to write it.

Writing a unit test to verify the functionality becomes pretty straightforward. e.g.

class TestTodo(unittest.TestCase):

    def test_run(self):
        filename = 'test.csv'
        user_inputs = [('a', 1), ('b', 2)]
        csv_contents = 'contents'
        todos_to_csv = Mock(return_value=csv_contents)
        write_rows = Mock()

        run(filename, user_inputs, todos_to_csv, write_rows)

        todos_to_csv.assert_called_with(user_inputs)
        write_rows.assert_called_with(filename, csv_contents)

    def test_todos_to_csv(self):
        user_inputs = [('a', 1), ('b', 2)]
        expected_contents = 'a,1\nb,2\n'

        contents = todos_to_csv(user_inputs)
        self.assertEqual(contents, expected_contents)

Breaking out your code into different bits of functionality has a lot of benefits. Things become more organized and easier to understand. Simple functions are easier to test and less likely to have bugs. By controlling what your code depends on, you can isolate different bits of functionality.

link to project code

Leave a comment