Added if __name__ == "__main__": to allow importing as a module.