Square code encode and decode I am doing some python practice and wrote a square code encoder/decoder. As described in drill #4 of Programming Practice Problems, the characters of a message are written left-to-right on a square grid, and each line of the encoded message is read by going top to bottom.

I'd like some feedback on what I wrote to see if there is anything I could have done better:

import math


def encode(message):
    message = message.replace(" ", "")
    square = math.ceil(math.sqrt(len(message)))

    # Create the "grid"
    rows = []
    for row in range(square):
        pos = row * square
        code = message[pos:pos + square]
        rows.append(code)

    print("The grid:\n{}".format("\n".join(rows)))

    # Encode the message
    coded_msg = ""
    for col in range(square):
        for row in range(square):
            # The last column of the last row will throw an out of range error, so handle that
            if col &lt; len(rows[row]):
                coded_msg += rows[row][col]
        coded_msg += " "

    return coded_msg.rstrip()


def decode(message):
    square = math.sqrt(len(message))
    col_count = math.ceil(square)
    row_count = math.ceil(len(message) / col_count)

    cols = message.split(" ")

    decoded_msg = ""
    for row in range(row_count):
        for col in cols:
            if row &lt; len(col):
                decoded_msg += col[row]

    return decoded_msg


if __name__ == "__main__":

    print("Gimme a message:")
    in_message = input()

    coded_message = encode(in_message)
    print("--\nSquare message:\n{}".format(coded_message))
    decoded_message = decode(coded_message)
    #print("--\nDecoded message:\n{}".format(decoded_message))

    print("--\ninput: {}\noutput: {}".format(in_message, decoded_message))
Couple of pointers.

In general it's good practice to avoid putting anything other than a call to another function (such as main()) in the if __name__ == "__main__" function, just to aid in ease of unit testing.

Document strings for functions are also nice to have:

def foo():
    '''This function does foo-all, always returns None'''
    pass


A nice (ish) one liner for the decode function might be something like this:

def encode(message):
    message = message.replace(" ", "")
    length = math.ceil(math.sqrt(len(message)))
    ''.join(message[i::length] i in range(length))


Although I admit it needs a bit of refactoring as its a tad long for one line (EDIT adapted a very nice improvement from 200_success 's answer, which they explain well). It does demonstrate that you can skip building the grid though.