#!/usr/bin/env python2

import os, sys, re, traceback, getopt, time, urllib, string

Y2KCUTOFF=60
__version__ = "0.1"

def showVersion():
  print 'pyQ v'+__version__+', by Rimon Barr:'
  print 'Python Yahoo Quote fetching utility'

def showUsage():
  print
  showVersion()
  print
  print 'Usage: pyQ [-i] [start_date [end_date]] ticker [ticker...]'
  print '       rimdu -h | -v'
  print 
  print '  -h, -?, --help      display this help information'
  print '  -v, --version       display version'
  print '  -i, --stdin         tickers fed on stdin, one per line'
  print
  print '    date formats are yyyymmdd'
  print '    if enddate is omitted, it is assume to be the same as startdate'
  print '    if startdate is omitted, we use *current* stock tables'
  print '      and otherwise, use historical stock tables.'
  print '      (current stock tables will give previous close price before'
  print '       market closing time.)'
  print '    tickers are exactly what you would type at finance.yahoo.com'
  print '    output format: "ticker, date (yyyymmdd), open, high, low, close, vol"'
  print '  date formats are yyyymmdd'
  print '  tickers are exactly what you would type at finance.yahoo.com'
  print '  output format: ticker, date, open, high, low, close, volume'
  print
  print 'Send comments, suggestions and bug reports to <barr+pyq@cs.cornell.edu>.'
  print

def usageError():
  print 'rimdu: command syntax error'
  print 'Try `rimdu --help\' for more information.'

def isInt(i):
  try:
    int(i)
    return 1
  except:
    return 0

# convert yyyymmdd string to tuple (yyyy, mm, dd)
def parseDate(d):
  return (d[-6:-4], d[-4:-2], d[-2:])

def yy2yyyy(yy):
  global Y2KCUTOFF;
  yy=int(yy) % 100
  if yy<Y2KCUTOFF: return `yy+2000`
  else: return `yy+1900`

# convert month to number
MONTH2NUM = {'Jan': 1, 'Feb': 2, 'Mar': 3, 'Apr': 4, 'May': 5, 'Jun': 6,
  'Jul': 7, 'Aug': 8, 'Sep': 9, 'Oct': 10, 'Nov': 11, 'Dec': 12}
def dd_mmm_yy2yyyymmdd(d):
  global MONTH2NUM
  d=string.split(d, '-')
  day='%02d' % int(d[0])
  month='%02d' % MONTH2NUM[d[1]]
  year=yy2yyyy(d[2])
  return year+month+day

def splitLines(buf):
  lines=string.split(buf, '\n')
  lines=filter(lambda x:x, lines)
  def removeCarriage(s):
    if s[-1]=='\r': return s[:-1]
    else: return s
  lines=[removeCarriage(l) for l in lines]
  return lines

def getTicker(d1, d2, ticker):
  d1=parseDate(d1)
  d2=parseDate(d2)
  url='http://table.finance.yahoo.com/table.csv'
  query={ 's':ticker, 'g':'d', 'x':'.csv', 'q':'q', 'y':'0',
    'a':d1[1], 'b':d1[2], 'c':d1[0][-2:],
    'd':d2[1], 'e':d2[2], 'f':d2[0][-2:]
  }
  url=url+'?'+urllib.urlencode(query)
  f=urllib.urlopen(url)
  buf=f.read()
  lines=splitLines(buf)
  if re.match('no prices', lines[0], re.I): return
  lines=lines[1:]
  def processLine(l, t=ticker):
    l=string.split(l, ',')
    l[0]=dd_mmm_yy2yyyymmdd(l[0])
    l=[t]+l
    print string.join(l, ',')
  for l in lines:
    processLine(l)

def getTickers(d1, d2, tickers):
  for t in tickers:
    getTicker(d1, d2, t)

def getTickersNowChunk(tickers):
  url='http://finance.yahoo.com/d/quotes.csv';
  tickers=string.join(tickers)
  query={ 's':tickers, 'f':'sohgpv', 'e':'.csv' }
  url=url+'?'+urllib.urlencode(query)
  f=urllib.urlopen(url)
  buf=f.read()
  lines=splitLines(buf)
  def processLine(l):
    l=string.split(l, ',')
    l[0]=string.lower(l[0][1:-1])
    t=time.localtime()
    l.insert(1, '%4d%02d%02d' % (t[0], t[1], t[2]))
    print string.join(l, ',')
  for l in lines:
    processLine(l)

def getTickersNow(tickers):
  while tickers:
    getTickersNowChunk(tickers[:150])
    tickers=tickers[150:]

def main():
  # parse options
  try:
    opts, args = getopt.getopt(sys.argv[1:], 'hv?i',
      ['help', 'version', 'stdin'])
  except getopt.GetoptError: 
    usageError()
    return
  # process options
  stdin=0
  for o, a in opts:
    if o in ("-h", "--help", "-?"):
      showUsage()
      return
    if o in ("-v", "--version"):
      showVersion()
      return
    if o in ("-i", "--stdin"):
      stdin=1
  t=time.localtime()
  startdate='%4d%02d%02d' % (t[0], t[1], t[2])
  enddate=startdate
  today=1
  tickers=[]
  argpos=-1
  for a in args:
    argpos=argpos+1
    if argpos==0 and isInt(a):
      startdate=enddate=a
      today=0
      continue
    if argpos==1 and isInt(a):
      enddate=a
      if a=='0':
        enddate='%4d%02d%02d' % (t[0], t[1], t[2])
      continue
    tickers=tickers+[a]
  if stdin:
    tickers=tickers+splitLines(sys.stdin.read())
  if not len(tickers):
    showUsage()
    return
  if today:
    getTickersNow(tickers)
  else:
    getTickers(startdate, enddate, tickers)

try:
  if __name__=='__main__': main()
except KeyboardInterrupt:
  traceback.print_exc()
  print 'Break!'