Coverage for common/inputfunc.py: 16%

1""" 

2crate_anon/common/inputfunc.py 

3 

4=============================================================================== 

5 

6 Copyright (C) 2015, University of Cambridge, Department of Psychiatry. 

7 Created by Rudolf Cardinal (rnc1001@cam.ac.uk). 

8 

9 This file is part of CRATE. 

10 

11 CRATE is free software: you can redistribute it and/or modify 

12 it under the terms of the GNU General Public License as published by 

13 the Free Software Foundation, either version 3 of the License, or 

14 (at your option) any later version. 

15 

16 CRATE is distributed in the hope that it will be useful, 

17 but WITHOUT ANY WARRANTY; without even the implied warranty of 

18 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

19 GNU General Public License for more details. 

20 

21 You should have received a copy of the GNU General Public License 

22 along with CRATE. If not, see <https://www.gnu.org/licenses/>. 

23 

24=============================================================================== 

25 

26**Read from files or stdin.** 

27 

28""" 

29 

30import logging 

31from typing import Iterable, List 

32 

33from cardinal_pythonlib.file_io import smart_open 

34 

35log = logging.getLogger(__name__) 

36 

37 

38# ============================================================================= 

39# Input 

40# ============================================================================= 

41 

42 

43def gen_chunks_from_files( 

44 filenames: List[str], 

45 chunk_terminator_line: str = None, 

46 stdin_prompt: str = None, 

47) -> Iterable[str]: 

48 """ 

49 Iterates through filenames (also permitting '-' for stdin). 

50 Generates multi-line chunks, separated by a terminator. 

51 

52 Args: 

53 filenames: 

54 Filenames (or '-' for stdin). 

55 chunk_terminator_line: 

56 Single-line string used to separate chunks within a file. If this 

57 is an empty string, then enter a blank line to terminate. If it is 

58 ``None``, then lines are yielded one by one (multi-line input is 

59 not possible). 

60 stdin_prompt: 

61 Optional prompt to show (to the log) before each request. 

62 

63 Yields: 

64 str: 

65 Each chunk. 

66 

67 """ 

68 current_lines = [] # type: List[str] 

69 use_prompt = False 

70 

71 def thing_to_yield() -> str: 

72 nonlocal current_lines 

73 chunk = "\n".join(current_lines) 

74 current_lines.clear() 

75 return chunk 

76 

77 def prompt_if_necessary() -> None: 

78 if use_prompt: 

79 log.warning(stdin_prompt) 

80 

81 stdin_filename = "-" 

82 for filename in filenames: 

83 log.info(f"Reading from: {filename}") 

84 use_prompt = stdin_filename and stdin_prompt 

85 with smart_open(filename) as f: 

86 prompt_if_necessary() 

87 for line in f: 

88 line = line.rstrip("\n") # remove trailing newline 

89 if chunk_terminator_line is None: 

90 yield line 

91 prompt_if_necessary() 

92 elif line == chunk_terminator_line: 

93 yield thing_to_yield() 

94 prompt_if_necessary() 

95 else: 

96 current_lines.append(line) 

97 # End of file: yield any leftovers 

98 yield thing_to_yield() 

99 log.debug(f"Finished file: {filename}")