Refactor graph.py