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