1 #! /usr/bin/env python

     2 

     3 # This file contains a class and a main program that perform three

     4 # related (though complimentary) formatting operations on Python

     5 # programs.  When called as "pindent -c", it takes a valid Python

     6 # program as input and outputs a version augmented with block-closing

     7 # comments.  When called as "pindent -d", it assumes its input is a

     8 # Python program with block-closing comments and outputs a commentless

     9 # version.   When called as "pindent -r" it assumes its input is a

    10 # Python program with block-closing comments but with its indentation

    11 # messed up, and outputs a properly indented version.

    12 

    13 # A "block-closing comment" is a comment of the form '# end <keyword>'

    14 # where <keyword> is the keyword that opened the block.  If the

    15 # opening keyword is 'def' or 'class', the function or class name may

    16 # be repeated in the block-closing comment as well.  Here is an

    17 # example of a program fully augmented with block-closing comments:

    18 

    19 # def foobar(a, b):

    20 #    if a == b:

    21 #        a = a+1

    22 #    elif a < b:

    23 #        b = b-1

    24 #        if b > a: a = a-1

    25 #        # end if

    26 #    else:

    27 #        print 'oops!'

    28 #    # end if

    29 # # end def foobar

    30 

    31 # Note that only the last part of an if...elif...else... block needs a

    32 # block-closing comment; the same is true for other compound

    33 # statements (e.g. try...except).  Also note that "short-form" blocks

    34 # like the second 'if' in the example must be closed as well;

    35 # otherwise the 'else' in the example would be ambiguous (remember

    36 # that indentation is not significant when interpreting block-closing

    37 # comments).

    38 

    39 # The operations are idempotent (i.e. applied to their own output

    40 # they yield an identical result).  Running first "pindent -c" and

    41 # then "pindent -r" on a valid Python program produces a program that

    42 # is semantically identical to the input (though its indentation may

    43 # be different). Running "pindent -e" on that output produces a

    44 # program that only differs from the original in indentation.

    45 

    46 # Other options:

    47 # -s stepsize: set the indentation step size (default 8)

    48 # -t tabsize : set the number of spaces a tab character is worth (default 8)

    49 # -e         : expand TABs into spaces

    50 # file ...   : input file(s) (default standard input)

    51 # The results always go to standard output

    52 

    53 # Caveats:

    54 # - comments ending in a backslash will be mistaken for continued lines

    55 # - continuations using backslash are always left unchanged

    56 # - continuations inside parentheses are not extra indented by -r

    57 #   but must be indented for -c to work correctly (this breaks

    58 #   idempotency!)

    59 # - continued lines inside triple-quoted strings are totally garbled

    60 

    61 # Secret feature:

    62 # - On input, a block may also be closed with an "end statement" --

    63 #   this is a block-closing comment without the '#' sign.

    64 

    65 # Possible improvements:

    66 # - check syntax based on transitions in 'next' table

    67 # - better error reporting

    68 # - better error recovery

    69 # - check identifier after class/def

    70 

    71 # The following wishes need a more complete tokenization of the source:

    72 # - Don't get fooled by comments ending in backslash

    73 # - reindent continuation lines indicated by backslash

    74 # - handle continuation lines inside parentheses/braces/brackets

    75 # - handle triple quoted strings spanning lines

    76 # - realign comments

    77 # - optionally do much more thorough reformatting, a la C indent

    78 

    79 # Defaults

    80 STEPSIZE = 8

    81 TABSIZE = 8

    82 EXPANDTABS = 0

    83 

    84 import os

    85 import re

    86 import sys

    87 

    88 next = {}

    89 next['if'] = next['elif'] = 'elif', 'else', 'end'

    90 next['while'] = next['for'] = 'else', 'end'

    91 next['try'] = 'except', 'finally'

    92 next['except'] = 'except', 'else', 'end'

    93 next['else'] = next['finally'] = next['def'] = next['class'] = 'end'

    94 next['end'] = ()

    95 start = 'if', 'while', 'for', 'try', 'def', 'class'

    96 

    97 class PythonIndenter:

    98 

    99     def __init__(self, fpi = sys.stdin, fpo = sys.stdout,

   100                  indentsize = STEPSIZE, tabsize = TABSIZE, expandtabs = EXPANDTABS):

   101         self.fpi = fpi

   102         self.fpo = fpo

   103         self.indentsize = indentsize

   104         self.tabsize = tabsize

   105         self.lineno = 0

   106         self.expandtabs = expandtabs

   107         self._write = fpo.write

   108         self.kwprog = re.compile(

   109                 r'^\s*(?P<kw>[a-z]+)'

   110                 r'(\s+(?P<id>[a-zA-Z_]\w*))?'

   111                 r'[^\w]')

   112         self.endprog = re.compile(

   113                 r'^\s*#?\s*end\s+(?P<kw>[a-z]+)'

   114                 r'(\s+(?P<id>[a-zA-Z_]\w*))?'

   115                 r'[^\w]')

   116         self.wsprog = re.compile(r'^[ \t]*')

   117     # end def __init__

   118 

   119     def write(self, line):

   120         if self.expandtabs:

   121             self._write(line.expandtabs(self.tabsize))

   122         else:

   123             self._write(line)

   124         # end if

   125     # end def write

   126 

   127     def readline(self):

   128         line = self.fpi.readline()

   129         if line: self.lineno = self.lineno + 1

   130         # end if

   131         return line

   132     # end def readline

   133 

   134     def error(self, fmt, *args):

   135         if args: fmt = fmt % args

   136         # end if

   137         sys.stderr.write('Error at line %d: %s\n' % (self.lineno, fmt))

   138         self.write('### %s ###\n' % fmt)

   139     # end def error

   140 

   141     def getline(self):

   142         line = self.readline()

   143         while line[-2:] == '\\\n':

   144             line2 = self.readline()

   145             if not line2: break

   146             # end if

   147             line = line + line2

   148         # end while

   149         return line

   150     # end def getline

   151 

   152     def putline(self, line, indent = None):

   153         if indent is None:

   154             self.write(line)

   155             return

   156         # end if

   157         tabs, spaces = divmod(indent*self.indentsize, self.tabsize)

   158         i = 0

   159         m = self.wsprog.match(line)

   160         if m: i = m.end()

   161         # end if

   162         self.write('\t'*tabs + ' '*spaces + line[i:])

   163     # end def putline

   164 

   165     def reformat(self):

   166         stack = []

   167         while 1:

   168             line = self.getline()

   169             if not line: break      # EOF

   170             # end if

   171             m = self.endprog.match(line)

   172             if m:

   173                 kw = 'end'

   174                 kw2 = m.group('kw')

   175                 if not stack:

   176                     self.error('unexpected end')

   177                 elif stack[-1][0] != kw2:

   178                     self.error('unmatched end')

   179                 # end if

   180                 del stack[-1:]

   181                 self.putline(line, len(stack))

   182                 continue

   183             # end if

   184             m = self.kwprog.match(line)

   185             if m:

   186                 kw = m.group('kw')

   187                 if kw in start:

   188                     self.putline(line, len(stack))

   189                     stack.append((kw, kw))

   190                     continue

   191                 # end if

   192                 if next.has_key(kw) and stack:

   193                     self.putline(line, len(stack)-1)

   194                     kwa, kwb = stack[-1]

   195                     stack[-1] = kwa, kw

   196                     continue

   197                 # end if

   198             # end if

   199             self.putline(line, len(stack))

   200         # end while

   201         if stack:

   202             self.error('unterminated keywords')

   203             for kwa, kwb in stack:

   204                 self.write('\t%s\n' % kwa)

   205             # end for

   206         # end if

   207     # end def reformat

   208 

   209     def delete(self):

   210         begin_counter = 0

   211         end_counter = 0

   212         while 1:

   213             line = self.getline()

   214             if not line: break      # EOF

   215             # end if

   216             m = self.endprog.match(line)

   217             if m:

   218                 end_counter = end_counter + 1

   219                 continue

   220             # end if

   221             m = self.kwprog.match(line)

   222             if m:

   223                 kw = m.group('kw')

   224                 if kw in start:

   225                     begin_counter = begin_counter + 1

   226                 # end if

   227             # end if

   228             self.putline(line)

   229         # end while

   230         if begin_counter - end_counter < 0:

   231             sys.stderr.write('Warning: input contained more end tags than expected\n')

   232         elif begin_counter - end_counter > 0:

   233             sys.stderr.write('Warning: input contained less end tags than expected\n')

   234         # end if

   235     # end def delete

   236 

   237     def complete(self):

   238         self.indentsize = 1

   239         stack = []

   240         todo = []

   241         thisid = ''

   242         current, firstkw, lastkw, topid = 0, '', '', ''

   243         while 1:

   244             line = self.getline()

   245             i = 0

   246             m = self.wsprog.match(line)

   247             if m: i = m.end()

   248             # end if

   249             m = self.endprog.match(line)

   250             if m:

   251                 thiskw = 'end'

   252                 endkw = m.group('kw')

   253                 thisid = m.group('id')

   254             else:

   255                 m = self.kwprog.match(line)

   256                 if m:

   257                     thiskw = m.group('kw')

   258                     if not next.has_key(thiskw):

   259                         thiskw = ''

   260                     # end if

   261                     if thiskw in ('def', 'class'):

   262                         thisid = m.group('id')

   263                     else:

   264                         thisid = ''

   265                     # end if

   266                 elif line[i:i+1] in ('\n', '#'):

   267                     todo.append(line)

   268                     continue

   269                 else:

   270                     thiskw = ''

   271                 # end if

   272             # end if

   273             indent = len(line[:i].expandtabs(self.tabsize))

   274             while indent < current:

   275                 if firstkw:

   276                     if topid:

   277                         s = '# end %s %s\n' % (

   278                                 firstkw, topid)

   279                     else:

   280                         s = '# end %s\n' % firstkw

   281                     # end if

   282                     self.putline(s, current)

   283                     firstkw = lastkw = ''

   284                 # end if

   285                 current, firstkw, lastkw, topid = stack[-1]

   286                 del stack[-1]

   287             # end while

   288             if indent == current and firstkw:

   289                 if thiskw == 'end':

   290                     if endkw != firstkw:

   291                         self.error('mismatched end')

   292                     # end if

   293                     firstkw = lastkw = ''

   294                 elif not thiskw or thiskw in start:

   295                     if topid:

   296                         s = '# end %s %s\n' % (

   297                                 firstkw, topid)

   298                     else:

   299                         s = '# end %s\n' % firstkw

   300                     # end if

   301                     self.putline(s, current)

   302                     firstkw = lastkw = topid = ''

   303                 # end if

   304             # end if

   305             if indent > current:

   306                 stack.append((current, firstkw, lastkw, topid))

   307                 if thiskw and thiskw not in start:

   308                     # error

   309                     thiskw = ''

   310                 # end if

   311                 current, firstkw, lastkw, topid = \

   312                          indent, thiskw, thiskw, thisid

   313             # end if

   314             if thiskw:

   315                 if thiskw in start:

   316                     firstkw = lastkw = thiskw

   317                     topid = thisid

   318                 else:

   319                     lastkw = thiskw

   320                 # end if

   321             # end if

   322             for l in todo: self.write(l)

   323             # end for

   324             todo = []

   325             if not line: break

   326             # end if

   327             self.write(line)

   328         # end while

   329     # end def complete

   330 

   331 # end class PythonIndenter

   332 

   333 # Simplified user interface

   334 # - xxx_filter(input, output): read and write file objects

   335 # - xxx_string(s): take and return string object

   336 # - xxx_file(filename): process file in place, return true iff changed

   337 

   338 def complete_filter(input = sys.stdin, output = sys.stdout,

   339                     stepsize = STEPSIZE, tabsize = TABSIZE, expandtabs = EXPANDTABS):

   340     pi = PythonIndenter(input, output, stepsize, tabsize, expandtabs)

   341     pi.complete()

   342 # end def complete_filter

   343 

   344 def delete_filter(input= sys.stdin, output = sys.stdout,

   345                         stepsize = STEPSIZE, tabsize = TABSIZE, expandtabs = EXPANDTABS):

   346     pi = PythonIndenter(input, output, stepsize, tabsize, expandtabs)

   347     pi.delete()

   348 # end def delete_filter

   349 

   350 def reformat_filter(input = sys.stdin, output = sys.stdout,

   351                     stepsize = STEPSIZE, tabsize = TABSIZE, expandtabs = EXPANDTABS):

   352     pi = PythonIndenter(input, output, stepsize, tabsize, expandtabs)

   353     pi.reformat()

   354 # end def reformat_filter

   355 

   356 class StringReader:

   357     def __init__(self, buf):

   358         self.buf = buf

   359         self.pos = 0

   360         self.len = len(self.buf)

   361     # end def __init__

   362     def read(self, n = 0):

   363         if n <= 0:

   364             n = self.len - self.pos

   365         else:

   366             n = min(n, self.len - self.pos)

   367         # end if

   368         r = self.buf[self.pos : self.pos + n]

   369         self.pos = self.pos + n

   370         return r

   371     # end def read

   372     def readline(self):

   373         i = self.buf.find('\n', self.pos)

   374         return self.read(i + 1 - self.pos)

   375     # end def readline

   376     def readlines(self):

   377         lines = []

   378         line = self.readline()

   379         while line:

   380             lines.append(line)

   381             line = self.readline()

   382         # end while

   383         return lines

   384     # end def readlines

   385     # seek/tell etc. are left as an exercise for the reader

   386 # end class StringReader

   387 

   388 class StringWriter:

   389     def __init__(self):

   390         self.buf = ''

   391     # end def __init__

   392     def write(self, s):

   393         self.buf = self.buf + s

   394     # end def write

   395     def getvalue(self):

   396         return self.buf

   397     # end def getvalue

   398 # end class StringWriter

   399 

   400 def complete_string(source, stepsize = STEPSIZE, tabsize = TABSIZE, expandtabs = EXPANDTABS):

   401     input = StringReader(source)

   402     output = StringWriter()

   403     pi = PythonIndenter(input, output, stepsize, tabsize, expandtabs)

   404     pi.complete()

   405     return output.getvalue()

   406 # end def complete_string

   407 

   408 def delete_string(source, stepsize = STEPSIZE, tabsize = TABSIZE, expandtabs = EXPANDTABS):

   409     input = StringReader(source)

   410     output = StringWriter()

   411     pi = PythonIndenter(input, output, stepsize, tabsize, expandtabs)

   412     pi.delete()

   413     return output.getvalue()

   414 # end def delete_string

   415 

   416 def reformat_string(source, stepsize = STEPSIZE, tabsize = TABSIZE, expandtabs = EXPANDTABS):

   417     input = StringReader(source)

   418     output = StringWriter()

   419     pi = PythonIndenter(input, output, stepsize, tabsize, expandtabs)

   420     pi.reformat()

   421     return output.getvalue()

   422 # end def reformat_string

   423 

   424 def complete_file(filename, stepsize = STEPSIZE, tabsize = TABSIZE, expandtabs = EXPANDTABS):

   425     source = open(filename, 'r').read()

   426     result = complete_string(source, stepsize, tabsize, expandtabs)

   427     if source == result: return 0

   428     # end if

   429     import os

   430     try: os.rename(filename, filename + '~')

   431     except os.error: pass

   432     # end try

   433     f = open(filename, 'w')

   434     f.write(result)

   435     f.close()

   436     return 1

   437 # end def complete_file

   438 

   439 def delete_file(filename, stepsize = STEPSIZE, tabsize = TABSIZE, expandtabs = EXPANDTABS):

   440     source = open(filename, 'r').read()

   441     result = delete_string(source, stepsize, tabsize, expandtabs)

   442     if source == result: return 0

   443     # end if

   444     import os

   445     try: os.rename(filename, filename + '~')

   446     except os.error: pass

   447     # end try

   448     f = open(filename, 'w')

   449     f.write(result)

   450     f.close()

   451     return 1

   452 # end def delete_file

   453 

   454 def reformat_file(filename, stepsize = STEPSIZE, tabsize = TABSIZE, expandtabs = EXPANDTABS):

   455     source = open(filename, 'r').read()

   456     result = reformat_string(source, stepsize, tabsize, expandtabs)

   457     if source == result: return 0

   458     # end if

   459     import os

   460     try: os.rename(filename, filename + '~')

   461     except os.error: pass

   462     # end try

   463     f = open(filename, 'w')

   464     f.write(result)

   465     f.close()

   466     return 1

   467 # end def reformat_file

   468 

   469 # Test program when called as a script

   470 

   471 usage = """

   472 usage: pindent (-c|-d|-r) [-s stepsize] [-t tabsize] [-e] [file] ...

   473 -c         : complete a correctly indented program (add #end directives)

   474 -d         : delete #end directives

   475 -r         : reformat a completed program (use #end directives)

   476 -s stepsize: indentation step (default %(STEPSIZE)d)

   477 -t tabsize : the worth in spaces of a tab (default %(TABSIZE)d)

   478 -e         : expand TABs into spaces (defailt OFF)

   479 [file] ... : files are changed in place, with backups in file~

   480 If no files are specified or a single - is given,

   481 the program acts as a filter (reads stdin, writes stdout).

   482 """ % vars()

   483 

   484 def error_both(op1, op2):

   485     sys.stderr.write('Error: You can not specify both '+op1+' and -'+op2[0]+' at the same time\n')

   486     sys.stderr.write(usage)

   487     sys.exit(2)

   488 # end def error_both

   489 

   490 def test():

   491     import getopt

   492     try:

   493         opts, args = getopt.getopt(sys.argv[1:], 'cdrs:t:e')

   494     except getopt.error, msg:

   495         sys.stderr.write('Error: %s\n' % msg)

   496         sys.stderr.write(usage)

   497         sys.exit(2)

   498     # end try

   499     action = None

   500     stepsize = STEPSIZE

   501     tabsize = TABSIZE

   502     expandtabs = EXPANDTABS

   503     for o, a in opts:

   504         if o == '-c':

   505             if action: error_both(o, action)

   506             # end if

   507             action = 'complete'

   508         elif o == '-d':

   509             if action: error_both(o, action)

   510             # end if

   511             action = 'delete'

   512         elif o == '-r':

   513             if action: error_both(o, action)

   514             # end if

   515             action = 'reformat'

   516         elif o == '-s':

   517             stepsize = int(a)

   518         elif o == '-t':

   519             tabsize = int(a)

   520         elif o == '-e':

   521             expandtabs = 1

   522         # end if

   523     # end for

   524     if not action:

   525         sys.stderr.write(

   526                 'You must specify -c(omplete), -d(elete) or -r(eformat)\n')

   527         sys.stderr.write(usage)

   528         sys.exit(2)

   529     # end if

   530     if not args or args == ['-']:

   531         action = eval(action + '_filter')

   532         action(sys.stdin, sys.stdout, stepsize, tabsize, expandtabs)

   533     else:

   534         action = eval(action + '_file')

   535         for filename in args:

   536             action(filename, stepsize, tabsize, expandtabs)

   537         # end for

   538     # end if

   539 # end def test

   540 

   541 if __name__ == '__main__':

   542     test()

   543 # end if