Add the start_date feature

RoughInvestmentReturns.py

@@ -1,15 +1,34 @@
 from __future__ import division
 
 import math, numpy
+from datetime import date
+
+def get_months(start_date, end_date):
+    '''
+    Gets the monthes between two dates
+    '''
+    start_month, start_year = start_date.split(":")
+    end_month, end_year = end_date.split(":")
+    return (int(end_year)-int(start_year))*12 + (int(end_month)-int(start_month))
 
 invested = int(input("Invested: "))
 current = int(input("Current: "))
-months = int(input("Months: "))
+months = input("Enter either Number of Months or A Start Date (MM:YYYY): ")
+
+if ":" in months:
+    # Create an object of the datetime module
+    today = date.today()
+    # Set the format to be MM:YYYY
+    date = today.strftime("%m:%Y")
+    months = get_months(start_date= months, end_date= date)
+else:
+    months = int(months)
+
 
 values = []
 
 for i in range(months):
-	values.append(-(invested/months))
+    values.append(-(invested/months))
 
 values.append(current)