fileinput.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442
  1. """Helper class to quickly write a loop over all standard input files.
  2. Typical use is:
  3. import fileinput
  4. for line in fileinput.input(encoding="utf-8"):
  5. process(line)
  6. This iterates over the lines of all files listed in sys.argv[1:],
  7. defaulting to sys.stdin if the list is empty. If a filename is '-' it
  8. is also replaced by sys.stdin and the optional arguments mode and
  9. openhook are ignored. To specify an alternative list of filenames,
  10. pass it as the argument to input(). A single file name is also allowed.
  11. Functions filename(), lineno() return the filename and cumulative line
  12. number of the line that has just been read; filelineno() returns its
  13. line number in the current file; isfirstline() returns true iff the
  14. line just read is the first line of its file; isstdin() returns true
  15. iff the line was read from sys.stdin. Function nextfile() closes the
  16. current file so that the next iteration will read the first line from
  17. the next file (if any); lines not read from the file will not count
  18. towards the cumulative line count; the filename is not changed until
  19. after the first line of the next file has been read. Function close()
  20. closes the sequence.
  21. Before any lines have been read, filename() returns None and both line
  22. numbers are zero; nextfile() has no effect. After all lines have been
  23. read, filename() and the line number functions return the values
  24. pertaining to the last line read; nextfile() has no effect.
  25. All files are opened in text mode by default, you can override this by
  26. setting the mode parameter to input() or FileInput.__init__().
  27. If an I/O error occurs during opening or reading a file, the OSError
  28. exception is raised.
  29. If sys.stdin is used more than once, the second and further use will
  30. return no lines, except perhaps for interactive use, or if it has been
  31. explicitly reset (e.g. using sys.stdin.seek(0)).
  32. Empty files are opened and immediately closed; the only time their
  33. presence in the list of filenames is noticeable at all is when the
  34. last file opened is empty.
  35. It is possible that the last line of a file doesn't end in a newline
  36. character; otherwise lines are returned including the trailing
  37. newline.
  38. Class FileInput is the implementation; its methods filename(),
  39. lineno(), fileline(), isfirstline(), isstdin(), nextfile() and close()
  40. correspond to the functions in the module. In addition it has a
  41. readline() method which returns the next input line, and a
  42. __getitem__() method which implements the sequence behavior. The
  43. sequence must be accessed in strictly sequential order; sequence
  44. access and readline() cannot be mixed.
  45. Optional in-place filtering: if the keyword argument inplace=1 is
  46. passed to input() or to the FileInput constructor, the file is moved
  47. to a backup file and standard output is directed to the input file.
  48. This makes it possible to write a filter that rewrites its input file
  49. in place. If the keyword argument backup=".<some extension>" is also
  50. given, it specifies the extension for the backup file, and the backup
  51. file remains around; by default, the extension is ".bak" and it is
  52. deleted when the output file is closed. In-place filtering is
  53. disabled when standard input is read. XXX The current implementation
  54. does not work for MS-DOS 8+3 filesystems.
  55. """
  56. import io
  57. import sys, os
  58. from types import GenericAlias
  59. __all__ = ["input", "close", "nextfile", "filename", "lineno", "filelineno",
  60. "fileno", "isfirstline", "isstdin", "FileInput", "hook_compressed",
  61. "hook_encoded"]
  62. _state = None
  63. def input(files=None, inplace=False, backup="", *, mode="r", openhook=None,
  64. encoding=None, errors=None):
  65. """Return an instance of the FileInput class, which can be iterated.
  66. The parameters are passed to the constructor of the FileInput class.
  67. The returned instance, in addition to being an iterator,
  68. keeps global state for the functions of this module,.
  69. """
  70. global _state
  71. if _state and _state._file:
  72. raise RuntimeError("input() already active")
  73. _state = FileInput(files, inplace, backup, mode=mode, openhook=openhook,
  74. encoding=encoding, errors=errors)
  75. return _state
  76. def close():
  77. """Close the sequence."""
  78. global _state
  79. state = _state
  80. _state = None
  81. if state:
  82. state.close()
  83. def nextfile():
  84. """
  85. Close the current file so that the next iteration will read the first
  86. line from the next file (if any); lines not read from the file will
  87. not count towards the cumulative line count. The filename is not
  88. changed until after the first line of the next file has been read.
  89. Before the first line has been read, this function has no effect;
  90. it cannot be used to skip the first file. After the last line of the
  91. last file has been read, this function has no effect.
  92. """
  93. if not _state:
  94. raise RuntimeError("no active input()")
  95. return _state.nextfile()
  96. def filename():
  97. """
  98. Return the name of the file currently being read.
  99. Before the first line has been read, returns None.
  100. """
  101. if not _state:
  102. raise RuntimeError("no active input()")
  103. return _state.filename()
  104. def lineno():
  105. """
  106. Return the cumulative line number of the line that has just been read.
  107. Before the first line has been read, returns 0. After the last line
  108. of the last file has been read, returns the line number of that line.
  109. """
  110. if not _state:
  111. raise RuntimeError("no active input()")
  112. return _state.lineno()
  113. def filelineno():
  114. """
  115. Return the line number in the current file. Before the first line
  116. has been read, returns 0. After the last line of the last file has
  117. been read, returns the line number of that line within the file.
  118. """
  119. if not _state:
  120. raise RuntimeError("no active input()")
  121. return _state.filelineno()
  122. def fileno():
  123. """
  124. Return the file number of the current file. When no file is currently
  125. opened, returns -1.
  126. """
  127. if not _state:
  128. raise RuntimeError("no active input()")
  129. return _state.fileno()
  130. def isfirstline():
  131. """
  132. Returns true the line just read is the first line of its file,
  133. otherwise returns false.
  134. """
  135. if not _state:
  136. raise RuntimeError("no active input()")
  137. return _state.isfirstline()
  138. def isstdin():
  139. """
  140. Returns true if the last line was read from sys.stdin,
  141. otherwise returns false.
  142. """
  143. if not _state:
  144. raise RuntimeError("no active input()")
  145. return _state.isstdin()
  146. class FileInput:
  147. """FileInput([files[, inplace[, backup]]], *, mode=None, openhook=None)
  148. Class FileInput is the implementation of the module; its methods
  149. filename(), lineno(), fileline(), isfirstline(), isstdin(), fileno(),
  150. nextfile() and close() correspond to the functions of the same name
  151. in the module.
  152. In addition it has a readline() method which returns the next
  153. input line, and a __getitem__() method which implements the
  154. sequence behavior. The sequence must be accessed in strictly
  155. sequential order; random access and readline() cannot be mixed.
  156. """
  157. def __init__(self, files=None, inplace=False, backup="", *,
  158. mode="r", openhook=None, encoding=None, errors=None):
  159. if isinstance(files, str):
  160. files = (files,)
  161. elif isinstance(files, os.PathLike):
  162. files = (os.fspath(files), )
  163. else:
  164. if files is None:
  165. files = sys.argv[1:]
  166. if not files:
  167. files = ('-',)
  168. else:
  169. files = tuple(files)
  170. self._files = files
  171. self._inplace = inplace
  172. self._backup = backup
  173. self._savestdout = None
  174. self._output = None
  175. self._filename = None
  176. self._startlineno = 0
  177. self._filelineno = 0
  178. self._file = None
  179. self._isstdin = False
  180. self._backupfilename = None
  181. self._encoding = encoding
  182. self._errors = errors
  183. # We can not use io.text_encoding() here because old openhook doesn't
  184. # take encoding parameter.
  185. if (sys.flags.warn_default_encoding and
  186. "b" not in mode and encoding is None and openhook is None):
  187. import warnings
  188. warnings.warn("'encoding' argument not specified.",
  189. EncodingWarning, 2)
  190. # restrict mode argument to reading modes
  191. if mode not in ('r', 'rb'):
  192. raise ValueError("FileInput opening mode must be 'r' or 'rb'")
  193. self._mode = mode
  194. self._write_mode = mode.replace('r', 'w')
  195. if openhook:
  196. if inplace:
  197. raise ValueError("FileInput cannot use an opening hook in inplace mode")
  198. if not callable(openhook):
  199. raise ValueError("FileInput openhook must be callable")
  200. self._openhook = openhook
  201. def __del__(self):
  202. self.close()
  203. def close(self):
  204. try:
  205. self.nextfile()
  206. finally:
  207. self._files = ()
  208. def __enter__(self):
  209. return self
  210. def __exit__(self, type, value, traceback):
  211. self.close()
  212. def __iter__(self):
  213. return self
  214. def __next__(self):
  215. while True:
  216. line = self._readline()
  217. if line:
  218. self._filelineno += 1
  219. return line
  220. if not self._file:
  221. raise StopIteration
  222. self.nextfile()
  223. # repeat with next file
  224. def nextfile(self):
  225. savestdout = self._savestdout
  226. self._savestdout = None
  227. if savestdout:
  228. sys.stdout = savestdout
  229. output = self._output
  230. self._output = None
  231. try:
  232. if output:
  233. output.close()
  234. finally:
  235. file = self._file
  236. self._file = None
  237. try:
  238. del self._readline # restore FileInput._readline
  239. except AttributeError:
  240. pass
  241. try:
  242. if file and not self._isstdin:
  243. file.close()
  244. finally:
  245. backupfilename = self._backupfilename
  246. self._backupfilename = None
  247. if backupfilename and not self._backup:
  248. try: os.unlink(backupfilename)
  249. except OSError: pass
  250. self._isstdin = False
  251. def readline(self):
  252. while True:
  253. line = self._readline()
  254. if line:
  255. self._filelineno += 1
  256. return line
  257. if not self._file:
  258. return line
  259. self.nextfile()
  260. # repeat with next file
  261. def _readline(self):
  262. if not self._files:
  263. if 'b' in self._mode:
  264. return b''
  265. else:
  266. return ''
  267. self._filename = self._files[0]
  268. self._files = self._files[1:]
  269. self._startlineno = self.lineno()
  270. self._filelineno = 0
  271. self._file = None
  272. self._isstdin = False
  273. self._backupfilename = 0
  274. # EncodingWarning is emitted in __init__() already
  275. if "b" not in self._mode:
  276. encoding = self._encoding or "locale"
  277. else:
  278. encoding = None
  279. if self._filename == '-':
  280. self._filename = '<stdin>'
  281. if 'b' in self._mode:
  282. self._file = getattr(sys.stdin, 'buffer', sys.stdin)
  283. else:
  284. self._file = sys.stdin
  285. self._isstdin = True
  286. else:
  287. if self._inplace:
  288. self._backupfilename = (
  289. os.fspath(self._filename) + (self._backup or ".bak"))
  290. try:
  291. os.unlink(self._backupfilename)
  292. except OSError:
  293. pass
  294. # The next few lines may raise OSError
  295. os.rename(self._filename, self._backupfilename)
  296. self._file = open(self._backupfilename, self._mode,
  297. encoding=encoding, errors=self._errors)
  298. try:
  299. perm = os.fstat(self._file.fileno()).st_mode
  300. except OSError:
  301. self._output = open(self._filename, self._write_mode,
  302. encoding=encoding, errors=self._errors)
  303. else:
  304. mode = os.O_CREAT | os.O_WRONLY | os.O_TRUNC
  305. if hasattr(os, 'O_BINARY'):
  306. mode |= os.O_BINARY
  307. fd = os.open(self._filename, mode, perm)
  308. self._output = os.fdopen(fd, self._write_mode,
  309. encoding=encoding, errors=self._errors)
  310. try:
  311. os.chmod(self._filename, perm)
  312. except OSError:
  313. pass
  314. self._savestdout = sys.stdout
  315. sys.stdout = self._output
  316. else:
  317. # This may raise OSError
  318. if self._openhook:
  319. # Custom hooks made previous to Python 3.10 didn't have
  320. # encoding argument
  321. if self._encoding is None:
  322. self._file = self._openhook(self._filename, self._mode)
  323. else:
  324. self._file = self._openhook(
  325. self._filename, self._mode, encoding=self._encoding, errors=self._errors)
  326. else:
  327. self._file = open(self._filename, self._mode, encoding=encoding, errors=self._errors)
  328. self._readline = self._file.readline # hide FileInput._readline
  329. return self._readline()
  330. def filename(self):
  331. return self._filename
  332. def lineno(self):
  333. return self._startlineno + self._filelineno
  334. def filelineno(self):
  335. return self._filelineno
  336. def fileno(self):
  337. if self._file:
  338. try:
  339. return self._file.fileno()
  340. except ValueError:
  341. return -1
  342. else:
  343. return -1
  344. def isfirstline(self):
  345. return self._filelineno == 1
  346. def isstdin(self):
  347. return self._isstdin
  348. __class_getitem__ = classmethod(GenericAlias)
  349. def hook_compressed(filename, mode, *, encoding=None, errors=None):
  350. if encoding is None and "b" not in mode: # EncodingWarning is emitted in FileInput() already.
  351. encoding = "locale"
  352. ext = os.path.splitext(filename)[1]
  353. if ext == '.gz':
  354. import gzip
  355. stream = gzip.open(filename, mode)
  356. elif ext == '.bz2':
  357. import bz2
  358. stream = bz2.BZ2File(filename, mode)
  359. else:
  360. return open(filename, mode, encoding=encoding, errors=errors)
  361. # gzip and bz2 are binary mode by default.
  362. if "b" not in mode:
  363. stream = io.TextIOWrapper(stream, encoding=encoding, errors=errors)
  364. return stream
  365. def hook_encoded(encoding, errors=None):
  366. def openhook(filename, mode):
  367. return open(filename, mode, encoding=encoding, errors=errors)
  368. return openhook
  369. def _test():
  370. import getopt
  371. inplace = False
  372. backup = False
  373. opts, args = getopt.getopt(sys.argv[1:], "ib:")
  374. for o, a in opts:
  375. if o == '-i': inplace = True
  376. if o == '-b': backup = a
  377. for line in input(args, inplace=inplace, backup=backup):
  378. if line[-1:] == '\n': line = line[:-1]
  379. if line[-1:] == '\r': line = line[:-1]
  380. print("%d: %s[%d]%s %s" % (lineno(), filename(), filelineno(),
  381. isfirstline() and "*" or "", line))
  382. print("%d: %s[%d]" % (lineno(), filename(), filelineno()))
  383. if __name__ == '__main__':
  384. _test()