1#!/usr/bin/env python3
2
3# Change the #! line (shebang) occurring in Python scripts.  The new interpreter
4# pathname must be given with a -i option.
5#
6# Command line arguments are files or directories to be processed.
7# Directories are searched recursively for files whose name looks
8# like a python module.
9# Symbolic links are always ignored (except as explicit directory
10# arguments).
11# The original file is kept as a back-up (with a "~" attached to its name),
12# -n flag can be used to disable this.
13
14# Sometimes you may find shebangs with flags such as `#! /usr/bin/env python -si`.
15# Normally, pathfix overwrites the entire line, including the flags.
16# To change interpreter and keep flags from the original shebang line, use -k.
17# If you want to keep flags and add to them one single literal flag, use option -a.
18
19
20# Undoubtedly you can do this using find and sed or perl, but this is
21# a nice example of Python code that recurses down a directory tree
22# and uses regular expressions.  Also note several subtleties like
23# preserving the file's mode and avoiding to even write a temp file
24# when no changes are needed for a file.
25#
26# NB: by changing only the function fixfile() you can turn this
27# into a program for a different change to Python programs...
28
29import sys
30import re
31import os
32from stat import *
33import getopt
34
35err = sys.stderr.write
36dbg = err
37rep = sys.stdout.write
38
39new_interpreter = None
40preserve_timestamps = False
41create_backup = True
42keep_flags = False
43add_flags = b''
44
45
46def main():
47    global new_interpreter
48    global preserve_timestamps
49    global create_backup
50    global keep_flags
51    global add_flags
52
53    usage = ('usage: %s -i /interpreter -p -n -k -a file-or-directory ...\n' %
54             sys.argv[0])
55    try:
56        opts, args = getopt.getopt(sys.argv[1:], 'i:a:kpn')
57    except getopt.error as msg:
58        err(str(msg) + '\n')
59        err(usage)
60        sys.exit(2)
61    for o, a in opts:
62        if o == '-i':
63            new_interpreter = a.encode()
64        if o == '-p':
65            preserve_timestamps = True
66        if o == '-n':
67            create_backup = False
68        if o == '-k':
69            keep_flags = True
70        if o == '-a':
71            add_flags = a.encode()
72            if b' ' in add_flags:
73                err("-a option doesn't support whitespaces")
74                sys.exit(2)
75    if not new_interpreter or not new_interpreter.startswith(b'/') or \
76           not args:
77        err('-i option or file-or-directory missing\n')
78        err(usage)
79        sys.exit(2)
80    bad = 0
81    for arg in args:
82        if os.path.isdir(arg):
83            if recursedown(arg): bad = 1
84        elif os.path.islink(arg):
85            err(arg + ': will not process symbolic links\n')
86            bad = 1
87        else:
88            if fix(arg): bad = 1
89    sys.exit(bad)
90
91
92def ispython(name):
93    return name.endswith('.py')
94
95
96def recursedown(dirname):
97    dbg('recursedown(%r)\n' % (dirname,))
98    bad = 0
99    try:
100        names = os.listdir(dirname)
101    except OSError as msg:
102        err('%s: cannot list directory: %r\n' % (dirname, msg))
103        return 1
104    names.sort()
105    subdirs = []
106    for name in names:
107        if name in (os.curdir, os.pardir): continue
108        fullname = os.path.join(dirname, name)
109        if os.path.islink(fullname): pass
110        elif os.path.isdir(fullname):
111            subdirs.append(fullname)
112        elif ispython(name):
113            if fix(fullname): bad = 1
114    for fullname in subdirs:
115        if recursedown(fullname): bad = 1
116    return bad
117
118
119def fix(filename):
120##  dbg('fix(%r)\n' % (filename,))
121    try:
122        f = open(filename, 'rb')
123    except IOError as msg:
124        err('%s: cannot open: %r\n' % (filename, msg))
125        return 1
126    with f:
127        line = f.readline()
128        fixed = fixline(line)
129        if line == fixed:
130            rep(filename+': no change\n')
131            return
132        head, tail = os.path.split(filename)
133        tempname = os.path.join(head, '@' + tail)
134        try:
135            g = open(tempname, 'wb')
136        except IOError as msg:
137            err('%s: cannot create: %r\n' % (tempname, msg))
138            return 1
139        with g:
140            rep(filename + ': updating\n')
141            g.write(fixed)
142            BUFSIZE = 8*1024
143            while 1:
144                buf = f.read(BUFSIZE)
145                if not buf: break
146                g.write(buf)
147
148    # Finishing touch -- move files
149
150    mtime = None
151    atime = None
152    # First copy the file's mode to the temp file
153    try:
154        statbuf = os.stat(filename)
155        mtime = statbuf.st_mtime
156        atime = statbuf.st_atime
157        os.chmod(tempname, statbuf[ST_MODE] & 0o7777)
158    except OSError as msg:
159        err('%s: warning: chmod failed (%r)\n' % (tempname, msg))
160    # Then make a backup of the original file as filename~
161    if create_backup:
162        try:
163            os.rename(filename, filename + '~')
164        except OSError as msg:
165            err('%s: warning: backup failed (%r)\n' % (filename, msg))
166    else:
167        try:
168            os.remove(filename)
169        except OSError as msg:
170            err('%s: warning: removing failed (%r)\n' % (filename, msg))
171    # Now move the temp file to the original file
172    try:
173        os.rename(tempname, filename)
174    except OSError as msg:
175        err('%s: rename failed (%r)\n' % (filename, msg))
176        return 1
177    if preserve_timestamps:
178        if atime and mtime:
179            try:
180                os.utime(filename, (atime, mtime))
181            except OSError as msg:
182                err('%s: reset of timestamp failed (%r)\n' % (filename, msg))
183                return 1
184    # Return success
185    return 0
186
187
188def parse_shebang(shebangline):
189    shebangline = shebangline.rstrip(b'\n')
190    start = shebangline.find(b' -')
191    if start == -1:
192        return b''
193    return shebangline[start:]
194
195
196def populate_flags(shebangline):
197    old_flags = b''
198    if keep_flags:
199        old_flags = parse_shebang(shebangline)
200        if old_flags:
201            old_flags = old_flags[2:]
202    if not (old_flags or add_flags):
203        return b''
204    # On Linux, the entire string following the interpreter name
205    # is passed as a single argument to the interpreter.
206    # e.g. "#! /usr/bin/python3 -W Error -s" runs "/usr/bin/python3 "-W Error -s"
207    # so shebang should have single '-' where flags are given and
208    # flag might need argument for that reasons adding new flags is
209    # between '-' and original flags
210    # e.g. #! /usr/bin/python3 -sW Error
211    return b' -' + add_flags + old_flags
212
213
214def fixline(line):
215    if not line.startswith(b'#!'):
216        return line
217
218    if b"python" not in line:
219        return line
220
221    flags = populate_flags(line)
222    return b'#! ' + new_interpreter + flags + b'\n'
223
224
225if __name__ == '__main__':
226    main()
227