blob: b89d4cdb58c9040a87e97212c3c2f621caa97f98 [file] [log] [blame]
Michael Walsh7423c012016-10-04 10:27:21 -05001#!/usr/bin/env python
2
3r"""
Michael Walsh410b1782019-10-22 15:56:18 -05004This module provides valuable argument processing functions like gen_get_options and sprint_args.
Michael Walsh7423c012016-10-04 10:27:21 -05005"""
6
7import sys
Michael Walshc0e0ad42019-12-05 16:54:59 -06008import os
9import re
10try:
11 import psutil
12 psutil_imported = True
13except ImportError:
14 psutil_imported = False
George Keishing3b7115a2018-08-02 10:48:17 -050015try:
16 import __builtin__
17except ImportError:
18 import builtins as __builtin__
Michael Walsh7423c012016-10-04 10:27:21 -050019import atexit
20import signal
21import argparse
Michael Walsh3f70fc52020-03-27 12:04:24 -050022import textwrap as textwrap
Michael Walsh7423c012016-10-04 10:27:21 -050023
24import gen_print as gp
Michael Walsh69d58ae2018-06-01 15:18:57 -050025import gen_valid as gv
Michael Walshc0e0ad42019-12-05 16:54:59 -060026import gen_cmd as gc
27import gen_misc as gm
Michael Walsh7423c012016-10-04 10:27:21 -050028
Michael Walsh3f70fc52020-03-27 12:04:24 -050029
30class MultilineFormatter(argparse.HelpFormatter):
31 def _fill_text(self, text, width, indent):
32 r"""
33 Split text into formatted lines for every "%%n" encountered in the text and return the result.
34 """
35 lines = self._whitespace_matcher.sub(' ', text).strip().split('%n')
36 formatted_lines = \
37 [textwrap.fill(x, width, initial_indent=indent, subsequent_indent=indent) + '\n' for x in lines]
38 return ''.join(formatted_lines)
39
40
41class ArgumentDefaultsHelpMultilineFormatter(MultilineFormatter, argparse.ArgumentDefaultsHelpFormatter):
42 pass
43
44
Michael Walsh7423c012016-10-04 10:27:21 -050045default_string = ' The default value is "%(default)s".'
Michael Walshc0e0ad42019-12-05 16:54:59 -060046module = sys.modules["__main__"]
Michael Walsh7423c012016-10-04 10:27:21 -050047
48
Michael Walsh7423c012016-10-04 10:27:21 -050049def gen_get_options(parser,
50 stock_list=[]):
Michael Walsh7423c012016-10-04 10:27:21 -050051 r"""
Michael Walsh410b1782019-10-22 15:56:18 -050052 Parse the command line arguments using the parser object passed and return True/False (i.e. pass/fail).
53 However, if gv.exit_on_error is set, simply exit the program on failure. Also set the following built in
54 values:
Michael Walsh7423c012016-10-04 10:27:21 -050055
56 __builtin__.quiet This value is used by the qprint functions.
57 __builtin__.test_mode This value is used by command processing functions.
58 __builtin__.debug This value is used by the dprint functions.
59 __builtin__.arg_obj This value is used by print_program_header, etc.
60 __builtin__.parser This value is used by print_program_header, etc.
61
62 Description of arguments:
Michael Walsh410b1782019-10-22 15:56:18 -050063 parser A parser object. See argparse module documentation for details.
64 stock_list The caller can use this parameter to request certain stock parameters
65 offered by this function. For example, this function will define a
66 "quiet" option upon request. This includes stop help text and parm
67 checking. The stock_list is a list of tuples each of which consists of
68 an arg_name and a default value. Example: stock_list = [("test_mode",
69 0), ("quiet", 1), ("debug", 0)]
Michael Walsh7423c012016-10-04 10:27:21 -050070 """
71
72 # This is a list of stock parms that we support.
73 master_stock_list = ["quiet", "test_mode", "debug", "loglevel"]
74
75 # Process stock_list.
76 for ix in range(0, len(stock_list)):
77 if len(stock_list[ix]) < 1:
Michael Walsh69d58ae2018-06-01 15:18:57 -050078 error_message = "Programmer error - stock_list[" + str(ix) +\
79 "] is supposed to be a tuple containing at" +\
80 " least one element which is the name of" +\
81 " the desired stock parameter:\n" +\
82 gp.sprint_var(stock_list)
83 return gv.process_error_message(error_message)
Joy Onyerikwu004ad3c2018-06-11 16:29:56 -050084 if isinstance(stock_list[ix], tuple):
Michael Walsh7423c012016-10-04 10:27:21 -050085 arg_name = stock_list[ix][0]
86 default = stock_list[ix][1]
87 else:
88 arg_name = stock_list[ix]
89 default = None
90
91 if arg_name not in master_stock_list:
Michael Walsh69d58ae2018-06-01 15:18:57 -050092 error_message = "Programmer error - arg_name \"" + arg_name +\
93 "\" not found found in stock list:\n" +\
94 gp.sprint_var(master_stock_list)
95 return gv.process_error_message(error_message)
Michael Walsh7423c012016-10-04 10:27:21 -050096
97 if arg_name == "quiet":
98 if default is None:
99 default = 0
100 parser.add_argument(
101 '--quiet',
102 default=default,
103 type=int,
104 choices=[1, 0],
Joy Onyerikwu004ad3c2018-06-11 16:29:56 -0500105 help='If this parameter is set to "1", %(prog)s'
106 + ' will print only essential information, i.e. it will'
107 + ' not echo parameters, echo commands, print the total'
108 + ' run time, etc.' + default_string)
Michael Walsh7423c012016-10-04 10:27:21 -0500109 elif arg_name == "test_mode":
110 if default is None:
111 default = 0
112 parser.add_argument(
113 '--test_mode',
114 default=default,
115 type=int,
116 choices=[1, 0],
Joy Onyerikwu004ad3c2018-06-11 16:29:56 -0500117 help='This means that %(prog)s should go through all the'
118 + ' motions but not actually do anything substantial.'
119 + ' This is mainly to be used by the developer of'
120 + ' %(prog)s.' + default_string)
Michael Walsh7423c012016-10-04 10:27:21 -0500121 elif arg_name == "debug":
122 if default is None:
123 default = 0
124 parser.add_argument(
125 '--debug',
126 default=default,
127 type=int,
128 choices=[1, 0],
Joy Onyerikwu004ad3c2018-06-11 16:29:56 -0500129 help='If this parameter is set to "1", %(prog)s will print'
130 + ' additional debug information. This is mainly to be'
131 + ' used by the developer of %(prog)s.' + default_string)
Michael Walsh7423c012016-10-04 10:27:21 -0500132 elif arg_name == "loglevel":
133 if default is None:
134 default = "info"
135 parser.add_argument(
136 '--loglevel',
137 default=default,
138 type=str,
139 choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL',
140 'debug', 'info', 'warning', 'error', 'critical'],
Joy Onyerikwu004ad3c2018-06-11 16:29:56 -0500141 help='If this parameter is set to "1", %(prog)s will print'
142 + ' additional debug information. This is mainly to be'
143 + ' used by the developer of %(prog)s.' + default_string)
Michael Walsh7423c012016-10-04 10:27:21 -0500144
145 arg_obj = parser.parse_args()
146
147 __builtin__.quiet = 0
148 __builtin__.test_mode = 0
149 __builtin__.debug = 0
150 __builtin__.loglevel = 'WARNING'
151 for ix in range(0, len(stock_list)):
Joy Onyerikwu004ad3c2018-06-11 16:29:56 -0500152 if isinstance(stock_list[ix], tuple):
Michael Walsh7423c012016-10-04 10:27:21 -0500153 arg_name = stock_list[ix][0]
154 default = stock_list[ix][1]
155 else:
156 arg_name = stock_list[ix]
157 default = None
158 if arg_name == "quiet":
159 __builtin__.quiet = arg_obj.quiet
160 elif arg_name == "test_mode":
161 __builtin__.test_mode = arg_obj.test_mode
162 elif arg_name == "debug":
163 __builtin__.debug = arg_obj.debug
164 elif arg_name == "loglevel":
165 __builtin__.loglevel = arg_obj.loglevel
166
167 __builtin__.arg_obj = arg_obj
168 __builtin__.parser = parser
169
Michael Walsh410b1782019-10-22 15:56:18 -0500170 # For each command line parameter, create a corresponding global variable and assign it the appropriate
171 # value. For example, if the command line contained "--last_name='Smith', we'll create a global variable
172 # named "last_name" with the value "Smith".
Michael Walsh7423c012016-10-04 10:27:21 -0500173 module = sys.modules['__main__']
174 for key in arg_obj.__dict__:
175 setattr(module, key, getattr(__builtin__.arg_obj, key))
176
177 return True
178
Michael Walsh7423c012016-10-04 10:27:21 -0500179
Michael Walshc33ef372017-01-10 11:46:29 -0600180def set_pgm_arg(var_value,
181 var_name=None):
Michael Walshc33ef372017-01-10 11:46:29 -0600182 r"""
Michael Walsh410b1782019-10-22 15:56:18 -0500183 Set the value of the arg_obj.__dict__ entry named in var_name with the var_value provided. Also, set
184 corresponding global variable.
Michael Walshc33ef372017-01-10 11:46:29 -0600185
186 Description of arguments:
187 var_value The value to set in the variable.
Michael Walsh410b1782019-10-22 15:56:18 -0500188 var_name The name of the variable to set. This defaults to the name of the
189 variable used for var_value when calling this function.
Michael Walshc33ef372017-01-10 11:46:29 -0600190 """
191
192 if var_name is None:
193 var_name = gp.get_arg_name(None, 1, 2)
194
195 arg_obj.__dict__[var_name] = var_value
196 module = sys.modules['__main__']
197 setattr(module, var_name, var_value)
198 if var_name == "quiet":
199 __builtin__.quiet = var_value
200 elif var_name == "debug":
201 __builtin__.debug = var_value
202 elif var_name == "test_mode":
203 __builtin__.test_mode = var_value
204
Michael Walshc33ef372017-01-10 11:46:29 -0600205
Michael Walsh7423c012016-10-04 10:27:21 -0500206def sprint_args(arg_obj,
207 indent=0):
Michael Walsh7423c012016-10-04 10:27:21 -0500208 r"""
Michael Walsh410b1782019-10-22 15:56:18 -0500209 sprint_var all of the arguments found in arg_obj and return the result as a string.
Michael Walsh7423c012016-10-04 10:27:21 -0500210
211 Description of arguments:
Michael Walsh410b1782019-10-22 15:56:18 -0500212 arg_obj An argument object such as is returned by the argparse parse_args()
213 method.
214 indent The number of spaces to indent each line of output.
Michael Walsh7423c012016-10-04 10:27:21 -0500215 """
216
Michael Walsh0d5f96a2019-05-20 10:09:57 -0500217 col1_width = gp.dft_col1_width + indent
Michael Walshbec416d2016-11-10 08:54:52 -0600218
Michael Walsh7423c012016-10-04 10:27:21 -0500219 buffer = ""
Michael Walsh7423c012016-10-04 10:27:21 -0500220 for key in arg_obj.__dict__:
Michael Walshbec416d2016-11-10 08:54:52 -0600221 buffer += gp.sprint_varx(key, getattr(arg_obj, key), 0, indent,
Michael Walsh0d5f96a2019-05-20 10:09:57 -0500222 col1_width)
Michael Walsh7423c012016-10-04 10:27:21 -0500223 return buffer
224
Michael Walsh7423c012016-10-04 10:27:21 -0500225
Michael Walshd339abf2020-02-18 11:28:02 -0600226def sync_args():
227 r"""
228 Synchronize the argument values to match their corresponding global variable values.
229
230 The user's validate_parms() function may manipulate global variables that correspond to program
231 arguments. After validate_parms() is called, sync_args is called to set the altered values back into the
232 arg_obj. This will ensure that the print-out of program arguments reflects the updated values.
233
234 Example:
235
236 def validate_parms():
237
238 # Set a default value for dir_path argument.
239 dir_path = gm.add_trailing_slash(gm.dft(dir_path, os.getcwd()))
240 """
241 module = sys.modules['__main__']
242 for key in arg_obj.__dict__:
243 arg_obj.__dict__[key] = getattr(module, key)
244
245
Michael Walshc0e0ad42019-12-05 16:54:59 -0600246term_options = None
Michael Walsh5328f382019-09-13 14:18:55 -0500247
248
Michael Walshc0e0ad42019-12-05 16:54:59 -0600249def set_term_options(**kwargs):
250 r"""
251 Set the global term_options.
252
253 If the global term_options is not None, gen_exit_function() will call terminate_descendants().
254
255 Description of arguments():
256 kwargs Supported keyword options follow:
257 term_requests Requests to terminate specified descendants of this program. The
258 following values for term_requests are supported:
259 children Terminate the direct children of this program.
260 descendants Terminate all descendants of this program.
261 <dictionary> A dictionary with support for the following keys:
262 pgm_names A list of program names which will be used to identify which descendant
263 processes should be terminated.
264 """
265
266 global term_options
267 # Validation:
268 arg_names = list(kwargs.keys())
269 gv.valid_list(arg_names, ['term_requests'])
270 if type(kwargs['term_requests']) is dict:
271 keys = list(kwargs['term_requests'].keys())
272 gv.valid_list(keys, ['pgm_names'])
273 else:
274 gv.valid_value(kwargs['term_requests'], ['children', 'descendants'])
275 term_options = kwargs
276
277
278if psutil_imported:
279 def match_process_by_pgm_name(process, pgm_name):
280 r"""
281 Return True or False to indicate whether the process matches the program name.
282
283 Description of argument(s):
284 process A psutil process object such as the one returned by psutil.Process().
285 pgm_name The name of a program to look for in the cmdline field of the process
286 object.
287 """
288
289 # This function will examine elements 0 and 1 of the cmdline field of the process object. The
290 # following examples will illustrate the reasons for this:
291
292 # Example 1: Suppose a process was started like this:
293
294 # shell_cmd('python_pgm_template --quiet=0', fork=1)
295
296 # And then this function is called as follows:
297
298 # match_process_by_pgm_name(process, "python_pgm_template")
299
300 # The process object might contain the following for its cmdline field:
301
302 # cmdline:
303 # [0]: /usr/bin/python
304 # [1]: /my_path/python_pgm_template
305 # [2]: --quiet=0
306
307 # Because "python_pgm_template" is a python program, the python interpreter (e.g. "/usr/bin/python")
308 # will appear in entry 0 of cmdline and the python_pgm_template will appear in entry 1 (with a
309 # qualifying dir path).
310
311 # Example 2: Suppose a process was started like this:
312
313 # shell_cmd('sleep 5', fork=1)
314
315 # And then this function is called as follows:
316
317 # match_process_by_pgm_name(process, "sleep")
318
319 # The process object might contain the following for its cmdline field:
320
321 # cmdline:
322 # [0]: sleep
323 # [1]: 5
324
325 # Because "sleep" is a compiled executable, it will appear in entry 0.
326
327 optional_dir_path_regex = "(.*/)?"
328 cmdline = process.as_dict()['cmdline']
329 return re.match(optional_dir_path_regex + pgm_name + '( |$)', cmdline[0]) \
330 or re.match(optional_dir_path_regex + pgm_name + '( |$)', cmdline[1])
331
332 def select_processes_by_pgm_name(processes, pgm_name):
333 r"""
334 Select the processes that match pgm_name and return the result as a list of process objects.
335
336 Description of argument(s):
337 processes A list of psutil process objects such as the one returned by
338 psutil.Process().
339 pgm_name The name of a program to look for in the cmdline field of each process
340 object.
341 """
342
343 return [process for process in processes if match_process_by_pgm_name(process, pgm_name)]
344
345 def sprint_process_report(pids):
346 r"""
347 Create a process report for the given pids and return it as a string.
348
349 Description of argument(s):
350 pids A list of process IDs for processes to be included in the report.
351 """
352 report = "\n"
353 cmd_buf = "echo ; ps wwo user,pgrp,pid,ppid,lstart,cmd --forest " + ' '.join(pids)
354 report += gp.sprint_issuing(cmd_buf)
355 rc, outbuf = gc.shell_cmd(cmd_buf, quiet=1)
356 report += outbuf + "\n"
357
358 return report
359
360 def get_descendant_info(process=psutil.Process()):
361 r"""
362 Get info about the descendants of the given process and return as a tuple of descendants,
363 descendant_pids and process_report.
364
365 descendants will be a list of process objects. descendant_pids will be a list of pids (in str form)
366 and process_report will be a report produced by a call to sprint_process_report().
367
368 Description of argument(s):
369 process A psutil process object such as the one returned by psutil.Process().
370 """
371 descendants = process.children(recursive=True)
372 descendant_pids = [str(process.pid) for process in descendants]
373 if descendants:
374 process_report = sprint_process_report([str(process.pid)] + descendant_pids)
375 else:
376 process_report = ""
377 return descendants, descendant_pids, process_report
378
379 def terminate_descendants():
380 r"""
381 Terminate descendants of the current process according to the requirements layed out in global
382 term_options variable.
383
384 Note: If term_options is not null, gen_exit_function() will automatically call this function.
385
386 When this function gets called, descendant processes may be running and may be printing to the same
387 stdout stream being used by this process. If this function writes directly to stdout, its output can
388 be interspersed with any output generated by descendant processes. This makes it very difficult to
389 interpret the output. In order solve this problem, the activity of this process will be logged to a
390 temporary file. After descendant processes have been terminated successfully, the temporary file
391 will be printed to stdout and then deleted. However, if this function should fail to complete (i.e.
392 get hung waiting for descendants to terminate gracefully), the temporary file will not be deleted and
393 can be used by the developer for debugging. If no descendant processes are found, this function will
394 return before creating the temporary file.
395
396 Note that a general principal being observed here is that each process is responsible for the
397 children it produces.
398 """
399
400 message = "\n" + gp.sprint_dashes(width=120) \
401 + gp.sprint_executing() + "\n"
402
403 current_process = psutil.Process()
404
405 descendants, descendant_pids, process_report = get_descendant_info(current_process)
406 if not descendants:
407 # If there are no descendants, then we have nothing to do.
408 return
409
410 terminate_descendants_temp_file_path = gm.create_temp_file_path()
411 gp.print_vars(terminate_descendants_temp_file_path)
412
413 message += gp.sprint_varx("pgm_name", gp.pgm_name) \
414 + gp.sprint_vars(term_options) \
415 + process_report
416
417 # Process the termination requests:
418 if term_options['term_requests'] == 'children':
419 term_processes = current_process.children(recursive=False)
420 term_pids = [str(process.pid) for process in term_processes]
421 elif term_options['term_requests'] == 'descendants':
422 term_processes = descendants
423 term_pids = descendant_pids
424 else:
425 # Process term requests by pgm_names.
426 term_processes = []
427 for pgm_name in term_options['term_requests']['pgm_names']:
428 term_processes.extend(select_processes_by_pgm_name(descendants, pgm_name))
429 term_pids = [str(process.pid) for process in term_processes]
430
431 message += gp.sprint_timen("Processes to be terminated:") \
432 + gp.sprint_var(term_pids)
433 for process in term_processes:
434 process.terminate()
435 message += gp.sprint_timen("Waiting on the following pids: " + ' '.join(descendant_pids))
436 gm.append_file(terminate_descendants_temp_file_path, message)
437 psutil.wait_procs(descendants)
438
439 # Checking after the fact to see whether any descendant processes are still alive. If so, a process
440 # report showing this will be included in the output.
441 descendants, descendant_pids, process_report = get_descendant_info(current_process)
442 if descendants:
443 message = "\n" + gp.sprint_timen("Not all of the processes terminated:") \
444 + process_report
445 gm.append_file(terminate_descendants_temp_file_path, message)
446
447 message = gp.sprint_dashes(width=120)
448 gm.append_file(terminate_descendants_temp_file_path, message)
449 gp.print_file(terminate_descendants_temp_file_path)
450 os.remove(terminate_descendants_temp_file_path)
451
452
453def gen_exit_function():
Michael Walsh5328f382019-09-13 14:18:55 -0500454 r"""
Michael Walsh410b1782019-10-22 15:56:18 -0500455 Execute whenever the program ends normally or with the signals that we catch (i.e. TERM, INT).
Michael Walsh5328f382019-09-13 14:18:55 -0500456 """
457
Michael Walshe15e7e92019-11-04 10:40:53 -0600458 # ignore_err influences the way shell_cmd processes errors. Since we're doing exit processing, we don't
459 # want to stop the program due to a shell_cmd failure.
460 ignore_err = 1
461
Michael Walshc0e0ad42019-12-05 16:54:59 -0600462 if psutil_imported and term_options:
463 terminate_descendants()
464
Michael Walsh5328f382019-09-13 14:18:55 -0500465 # Call the main module's exit_function if it is defined.
466 exit_function = getattr(module, "exit_function", None)
467 if exit_function:
Michael Walshc0e0ad42019-12-05 16:54:59 -0600468 exit_function()
Michael Walsh5328f382019-09-13 14:18:55 -0500469
470 gp.qprint_pgm_footer()
471
472
473def gen_signal_handler(signal_number,
474 frame):
475 r"""
Michael Walsh410b1782019-10-22 15:56:18 -0500476 Handle signals. Without a function to catch a SIGTERM or SIGINT, the program would terminate immediately
477 with return code 143 and without calling the exit_function.
Michael Walsh5328f382019-09-13 14:18:55 -0500478 """
479
Michael Walsh410b1782019-10-22 15:56:18 -0500480 # The convention is to set up exit_function with atexit.register() so there is no need to explicitly
481 # call exit_function from here.
Michael Walsh5328f382019-09-13 14:18:55 -0500482
Michael Walshc0e0ad42019-12-05 16:54:59 -0600483 gp.qprint_executing()
Michael Walsh5328f382019-09-13 14:18:55 -0500484
Michael Walsh410b1782019-10-22 15:56:18 -0500485 # Calling exit prevents control from returning to the code that was running when the signal was received.
Michael Walsh5328f382019-09-13 14:18:55 -0500486 exit(0)
487
488
Michael Walsh7423c012016-10-04 10:27:21 -0500489def gen_post_validation(exit_function=None,
490 signal_handler=None):
Michael Walsh7423c012016-10-04 10:27:21 -0500491 r"""
Michael Walsh410b1782019-10-22 15:56:18 -0500492 Do generic post-validation processing. By "post", we mean that this is to be called from a validation
493 function after the caller has done any validation desired. If the calling program passes exit_function
494 and signal_handler parms, this function will register them. In other words, it will make the
495 signal_handler functions get called for SIGINT and SIGTERM and will make the exit_function function run
496 prior to the termination of the program.
Michael Walsh7423c012016-10-04 10:27:21 -0500497
498 Description of arguments:
Michael Walsh410b1782019-10-22 15:56:18 -0500499 exit_function A function object pointing to the caller's exit function. This defaults
500 to this module's gen_exit_function.
501 signal_handler A function object pointing to the caller's signal_handler function. This
502 defaults to this module's gen_signal_handler.
Michael Walsh7423c012016-10-04 10:27:21 -0500503 """
504
Michael Walsh5328f382019-09-13 14:18:55 -0500505 # Get defaults.
506 exit_function = exit_function or gen_exit_function
507 signal_handler = signal_handler or gen_signal_handler
508
509 atexit.register(exit_function)
510 signal.signal(signal.SIGINT, signal_handler)
511 signal.signal(signal.SIGTERM, signal_handler)
512
513
514def gen_setup():
515 r"""
516 Do general setup for a program.
517 """
518
519 # Set exit_on_error for gen_valid functions.
520 gv.set_exit_on_error(True)
521
522 # Get main module variable values.
523 parser = getattr(module, "parser")
524 stock_list = getattr(module, "stock_list")
525 validate_parms = getattr(module, "validate_parms", None)
526
527 gen_get_options(parser, stock_list)
528
529 if validate_parms:
530 validate_parms()
Michael Walshd339abf2020-02-18 11:28:02 -0600531 sync_args()
Michael Walsh5328f382019-09-13 14:18:55 -0500532 gen_post_validation()
533
534 gp.qprint_pgm_header()