#
# Add line numbers to vb code
# Author: Chui Tey
#         teyc@cognoware.com
#
# Date:   05-Sep-2000
#         11-Nov-2000 Handles write locked files
#         12-Feb-2001 Restores original file MAC time
#         12-Feb-2001 Line number according to IDE line numbers
#                     handles both cls and frm files, untested on
#                     ctl files
#         24-Oct-2007 Handles ctl properly
#                     Line numbers inside select statements
#                     Makes sure that Functions and Subs aren't 
#                     line numbered, even if it's not expected to
#                     be present, as conditional defines can break
#                     the state machine

header = '<pre>\n'
footer = '</pre>\n'
header = ''
footer = ''

import re
import os
from string import strip
from string import split
from string import lower
from sys    import argv
from sys    import stdout

start_select = re.compile('^\s*Select Case')
end_select   = re.compile('^\s*End Select') # We are allowed to line number after the first Case
end_select   = re.compile('^\s*Case ')
start_sub = re.compile('^(\s*Public |\s*Private )?(Property|Function|Sub) ')
pattern_omit = re.compile("""^\d+\s|^\S+:|^\s*(\'|Attribute|Exit)|^\s*\#|^\s*Case""")
pattern_remove = re.compile('^\d+\s{1,4}')

def removeToFile(file):
  """
  Removes line numbers to a file of given name
  """

  # save the MAC time of the file
  atime, mtime, ctime = os.stat(file)[-3:]

  lines = open(file,'r').readlines()
  try:
    outfile = open(file,'w')
  except IOError:
    # file was locked
    return
  #
  # initializers
  state = ''
  line_continuation = 0
  outfile.write(header)
  for line in lines:
     # default
     new_line = line
     if strip(line) != '':
        if line_continuation:
           # previous line was a line continuation
           # don't modify the current line
           pass
        elif state != 'inside':
           #
           # we are outside a function or sub
           # find the start of a function or sub
           #
           match=start_sub.match(line)
           if match:
                state = 'inside'
                type = match.group(2)  # 'sub', 'function', 'property'

               # Some routines are all declared on the same line
               # e.g. Public Function Foo(): Foo = 123: End Function
            
                if line.find(' End '+ type) >= 0:
                    state = ''

        else:
           if (strip(line)[:4+len(type)])== ('End ' + type):
                state = ''
           elif pattern_remove.match(line):
                # this is a line number label
                # remove
                #
                new_line = pattern_remove.sub('', line)
        #
        # determine if current
        # line is continuation line 
        #
        line_continuation = is_continuation(line)
     else:
        line_continuation=0
     outfile.write(new_line)

  outfile.write(footer)
  outfile.close()

  # Restore original MAC time
  os.utime(file, (atime, mtime))

def addToFile(file):
  """
  Add lines numbers to file
  """
  #
  # default settings
  #
  line_no = 0
  line_no_incr = 1

  # save the MAC time of the file
  atime, mtime, ctime = os.stat(file)[-3:]

  #
  # initializers
  #
  lines = open(file,'r').readlines()
  outfilename = file # + ',v'  #xxx
  try:
    outfile = open(outfilename ,'w')
  except IOError:
    # file was locked
    return

  inside_select = 0
  state=''
  line_continuation=0

  # if this is a frm or cls file, then write the form
  # definition lines directly to the output file.
  # Form definitions are bounded by 'begin' and
  # 'end' statements
  #
  if lower(file[-3:]) in ('ctl', 'frm', 'cls'):

    num_linestoomit = 0
    num_begin = 0
    num_end   = 0

    for line in lines:
      if lower(strip(line)[:5]) == 'begin':
         num_begin = num_begin + 1
      elif lower(strip(line)[:3]) == 'end':
         num_end   = num_end   + 1

      num_linestoomit = num_linestoomit + 1
      outfile.write(line)

      if num_end > 0 and num_end == num_begin:
         lines = lines[num_linestoomit:]
         break

  for line in lines:

     # increment line number if we are not
     # in a 'Attribute' section
     #
     if strip(line)[0:9] != 'Attribute':
        line_no = line_no + line_no_incr

     # default
     new_line = line
     if strip(line) != '':
        if line_continuation:
           pass
        elif state != 'inside':
           match=start_sub.match(line)
           if match:
                state = 'inside'
                type = match.group(2)

               # Some routines are all declared on the same line
               # e.g. Public Function Foo(): Foo = 123: End Function
            
                if line.find(' End '+ type) >= 0:
                    state = ''

        elif start_sub.match(line):
           # due to #IF that can occur
           # we double check that we aren't
           # line numbering procedure calls
           pass

        else:
           #
           # Handle going into select case
           #
           if start_select.match(line):
                inside_select = 1
           if (strip(line)[:4+len(type)])== ('End ' + type):
                state = ''
           elif pattern_omit.match(line):
                pass # this is a line label
           else:
                if not inside_select:
                    new_line =  str(line_no) + '    ' + line
           if end_select.match(line):
                inside_select = 0

        #
        # determine if current
        # line is continuation line 
        #
        line_continuation = is_continuation(line)
     else:
        line_continuation=0
     outfile.write(new_line)

  outfile.write(footer)
  outfile.close()

  # Restore original MAC time
  os.utime(file, (atime, mtime))

  # Not needed, looks like readlines reads all into memory  
  if outfilename != file:
      os.unlink(file)
      os.rename(outfilename, file)

def is_continuation(line):
  """        
  determine if current
  line is continuation line 
  """        
  if line[-3:-1] == " _":
        return 1
  else:
        return 0

def main():
   if len(argv)==1:
        print "Usage: python LineNumber.py [-add|-remove] files [files]..."
        print "  file: wildcards allowed"
   else:
        cmdline = parse_args()
        #files = os.listdir('.')
        files = cmdline['file']
        vb_files = filter(lambda f:fileext(f) in ['ctl', 'cls', 'frm', 'bas'], files)
        for file in vb_files:
           if cmdline['add']:
               addToFile(file)
           else:
               removeToFile(file)
  
def fileext(filename):
   if filename[-4] =='.':
        return lower(filename[-3:])
   else:
        return None
   
def parse_args():
   from glob import glob
   result={'add':0, 'file': []}
   for arg in argv[1:]:
      if arg == '-add':
         result['add'] = 1
      else:
         result['file'].extend( glob(arg))
   return result

main()