Cleaned up code_gen.py a bit.