1import re
2import sys
3
4try:
5    import parser
6except ImportError:
7    parser = None
8
9d={}
10#  d is the dictionary of unittest changes, keyed to the old name
11#  used by unittest.
12#  d[old][0] is the new replacement function.
13#  d[old][1] is the operator you will substitute, or '' if there is none.
14#  d[old][2] is the possible number of arguments to the unittest
15#  function.
16
17# Old Unittest Name             new name         operator  # of args
18d['assertRaises']           = ('raises',               '', ['Any'])
19d['fail']                   = ('raise AssertionError', '', [0,1])
20d['assert_']                = ('assert',               '', [1,2])
21d['failIf']                 = ('assert not',           '', [1,2])
22d['assertEqual']            = ('assert',            ' ==', [2,3])
23d['failIfEqual']            = ('assert not',        ' ==', [2,3])
24d['assertIn']               = ('assert',            ' in', [2,3])
25d['assertNotIn']            = ('assert',            ' not in', [2,3])
26d['assertNotEqual']         = ('assert',            ' !=', [2,3])
27d['failUnlessEqual']        = ('assert',            ' ==', [2,3])
28d['assertAlmostEqual']      = ('assert round',      ' ==', [2,3,4])
29d['failIfAlmostEqual']      = ('assert not round',  ' ==', [2,3,4])
30d['assertNotAlmostEqual']   = ('assert round',      ' !=', [2,3,4])
31d['failUnlessAlmostEquals'] = ('assert round',      ' ==', [2,3,4])
32d['assertIsNone']           = ('assert',            ' is None', [1,2])
33d['assertIsNotNone']        = ('assert',            ' is not None', [1,2])
34d['assertGreater']          = ('assert',            ' >', [2,3])
35
36#  the list of synonyms
37d['failUnlessRaises']      = d['assertRaises']
38d['failUnless']            = d['assert_']
39d['assertTrue']            = d['assert_']
40d['assertFalse']           = d['failIf']
41d['assertEquals']          = d['assertEqual']
42d['assertNotEquals']       = d['assertNotEqual']
43d['assertAlmostEquals']    = d['assertAlmostEqual']
44d['assertNotAlmostEquals'] = d['assertNotAlmostEqual']
45
46# set up the regular expressions we will need
47leading_spaces = re.compile(r'^(\s*)') # this never fails
48
49pat = ''
50for k in d.keys():  # this complicated pattern to match all unittests
51    pat += '|' + r'^(\s*)' + 'self.' + k + r'\(' # \tself.whatever(
52
53old_names = re.compile(pat[1:])
54linesep='\n'        # nobody will really try to convert files not read
55                    # in text mode, will they?
56
57
58def blocksplitter(fp):
59    '''split a file into blocks that are headed by functions to rename'''
60
61    blocklist = []
62    blockstring = ''
63
64    for line in fp:
65        interesting = old_names.match(line)
66        if interesting :
67            if blockstring:
68                blocklist.append(blockstring)
69                blockstring = line # reset the block
70        else:
71            blockstring += line
72
73    blocklist.append(blockstring)
74    return blocklist
75
76def rewrite_utest(block):
77    '''rewrite every block to use the new utest functions'''
78
79    '''returns the rewritten unittest, unless it ran into problems,
80       in which case it just returns the block unchanged.
81    '''
82    utest = old_names.match(block)
83
84    if not utest:
85        return block
86
87    old = utest.group(0).lstrip()[5:-1] # the name we want to replace
88    new = d[old][0] # the name of the replacement function
89    op  = d[old][1] # the operator you will use , or '' if there is none.
90    possible_args = d[old][2]  # a list of the number of arguments the
91                               # unittest function could possibly take.
92
93    if possible_args == ['Any']: # just rename assertRaises & friends
94        return re.sub('self.'+old, new, block)
95
96    message_pos = possible_args[-1]
97    # the remaining unittests can have an optional message to print
98    # when they fail.  It is always the last argument to the function.
99
100    try:
101        indent, argl, trailer = decompose_unittest(old, block)
102
103    except SyntaxError: # but we couldn't parse it!
104        return block
105
106    argnum = len(argl)
107    if argnum not in possible_args:
108        # sanity check - this one isn't real either
109        return block
110
111    elif argnum == message_pos:
112        message = argl[-1]
113        argl = argl[:-1]
114    else:
115        message = None
116
117    if argnum is 0 or (argnum is 1 and argnum is message_pos): #unittest fail()
118        string = ''
119        if message:
120            message = ' ' + message
121
122    elif message_pos is 4:  # assertAlmostEqual & friends
123        try:
124            pos = argl[2].lstrip()
125        except IndexError:
126            pos = '7' # default if none is specified
127        string = '(%s -%s, %s)%s 0' % (argl[0], argl[1], pos, op )
128
129    elif old in ['assertIsNone', 'assertIsNotNone']:
130        string = ' ' + ''.join(argl) + op
131
132    else: # assert_, assertEquals and all the rest
133        string = ' ' + op.join(argl)
134
135    if message:
136        string = string + ',' + message
137
138    return indent + new + string + trailer
139
140def decompose_unittest(old, block):
141    '''decompose the block into its component parts'''
142
143    ''' returns indent, arglist, trailer
144        indent -- the indentation
145        arglist -- the arguments to the unittest function
146        trailer -- any extra junk after the closing paren, such as #commment
147    '''
148
149    indent = re.match(r'(\s*)', block).group()
150    pat = re.search('self.' + old + r'\(', block)
151
152    args, trailer = get_expr(block[pat.end():], ')')
153    arglist = break_args(args, [])
154
155    if arglist == ['']: # there weren't any
156        return indent, [], trailer
157
158    for i in range(len(arglist)):
159        try:
160            parser.expr(arglist[i].lstrip('\t '))
161        except SyntaxError:
162            if i == 0:
163                arglist[i] = '(' + arglist[i] + ')'
164            else:
165                arglist[i] = ' (' + arglist[i] + ')'
166
167    return indent, arglist, trailer
168
169def break_args(args, arglist):
170    '''recursively break a string into a list of arguments'''
171    try:
172        first, rest = get_expr(args, ',')
173        if not rest:
174            return arglist + [first]
175        else:
176            return [first] + break_args(rest, arglist)
177    except SyntaxError:
178        return arglist + [args]
179
180def get_expr(s, char):
181    '''split a string into an expression, and the rest of the string'''
182
183    pos=[]
184    for i in range(len(s)):
185        if s[i] == char:
186            pos.append(i)
187    if pos == []:
188        raise SyntaxError # we didn't find the expected char.  Ick.
189
190    for p in pos:
191        # make the python parser do the hard work of deciding which comma
192        # splits the string into two expressions
193        try:
194            parser.expr('(' + s[:p] + ')')
195            return s[:p], s[p+1:]
196        except SyntaxError: # It's not an expression yet
197            pass
198    raise SyntaxError       # We never found anything that worked.
199
200
201def main():
202    import sys
203    import py
204
205    usage = "usage: %prog [-s [filename ...] | [-i | -c filename ...]]"
206    optparser = py.std.optparse.OptionParser(usage)
207
208    def select_output (option, opt, value, optparser, **kw):
209        if hasattr(optparser, 'output'):
210            optparser.error(
211                'Cannot combine -s -i and -c options. Use one only.')
212        else:
213            optparser.output = kw['output']
214
215    optparser.add_option("-s", "--stdout", action="callback",
216                         callback=select_output,
217                         callback_kwargs={'output':'stdout'},
218                         help="send your output to stdout")
219
220    optparser.add_option("-i", "--inplace", action="callback",
221                         callback=select_output,
222                         callback_kwargs={'output':'inplace'},
223                         help="overwrite files in place")
224
225    optparser.add_option("-c", "--copy", action="callback",
226                         callback=select_output,
227                         callback_kwargs={'output':'copy'},
228                         help="copy files ... fn.py --> fn_cp.py")
229
230    options, args = optparser.parse_args()
231
232    output = getattr(optparser, 'output', 'stdout')
233
234    if output in ['inplace', 'copy'] and not args:
235        optparser.error(
236                '-i and -c option  require at least one filename')
237
238    if not args:
239        s = ''
240        for block in blocksplitter(sys.stdin):
241            s += rewrite_utest(block)
242        sys.stdout.write(s)
243
244    else:
245        for infilename in args: # no error checking to see if we can open, etc.
246            infile = file(infilename)
247            s = ''
248            for block in blocksplitter(infile):
249                s += rewrite_utest(block)
250            if output == 'inplace':
251                outfile = file(infilename, 'w+')
252            elif output == 'copy': # yes, just go clobber any existing .cp
253                outfile = file (infilename[:-3]+ '_cp.py', 'w+')
254            else:
255                outfile = sys.stdout
256
257            outfile.write(s)
258
259
260if __name__ == '__main__':
261    main()
262