diff --git a/counter.py b/counter.py index 1e2fb56..42e554c 100644 --- a/counter.py +++ b/counter.py @@ -1,8 +1,8 @@ """ A program that stores and updates a counter using a Python pickle file""" -from os.path import exists import sys -from pickle import dump, load +import os +import pickle def update_counter(file_name, reset=False): """ Updates a counter stored in the file 'file_name' @@ -29,7 +29,31 @@ def update_counter(file_name, reset=False): >>> update_counter('blah2.txt') 2 """ - pass + if os.path.exists(file_name) == False: + f = open(file_name, 'w') + counter = 1 + pickle.dump(counter, f) + f.close + return counter + + elif os.path.exists(file_name) == True: + if reset == True: + f = open(file_name, 'w') + counter = 1 + pickle.dump(counter, f) + f.close + return counter + + else: + f = open(file_name, 'r+') + counter = pickle.load(f) + counter += 1 + f.seek(0,0) + pickle.dump(counter, f) + f.close + return counter + + print counter if __name__ == '__main__': if len(sys.argv) < 2: