Module relax_io
[hide private]
[frames] | no frames]

Source Code for Module relax_io

   1  ############################################################################### 
   2  #                                                                             # 
   3  # Copyright (C) 2003-2012 Edward d'Auvergne                                   # 
   4  #                                                                             # 
   5  # This file is part of the program relax.                                     # 
   6  #                                                                             # 
   7  # relax is free software; you can redistribute it and/or modify               # 
   8  # it under the terms of the GNU General Public License as published by        # 
   9  # the Free Software Foundation; either version 2 of the License, or           # 
  10  # (at your option) any later version.                                         # 
  11  #                                                                             # 
  12  # relax is distributed in the hope that it will be useful,                    # 
  13  # but WITHOUT ANY WARRANTY; without even the implied warranty of              # 
  14  # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the               # 
  15  # GNU General Public License for more details.                                # 
  16  #                                                                             # 
  17  # You should have received a copy of the GNU General Public License           # 
  18  # along with relax; if not, write to the Free Software                        # 
  19  # Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA   # 
  20  #                                                                             # 
  21  ############################################################################### 
  22   
  23  # Module docstring. 
  24  """Module containing advanced IO functions for relax. 
  25   
  26  This includes IO redirection, automatic loading and writing of compressed files (both Gzip and BZ2 
  27  compression), reading and writing of files, processing of the contents of files, etc. 
  28  """ 
  29   
  30  # Dependency check module. 
  31  import dep_check 
  32   
  33  # Python module imports. 
  34  if dep_check.bz2_module: 
  35      from bz2 import BZ2File 
  36  if dep_check.gzip_module: 
  37      from gzip import GzipFile 
  38  if dep_check.devnull_import: 
  39      from os import devnull 
  40  from os import F_OK, X_OK, access, altsep, getenv, makedirs, pathsep, remove, sep 
  41  from os.path import expanduser, basename, splitext 
  42  from re import match, search 
  43  from string import split 
  44  import sys 
  45  from sys import stdin, stdout, stderr 
  46  from warnings import warn 
  47   
  48  # relax module imports. 
  49  import generic_fns 
  50  from generic_fns.mol_res_spin import generate_spin_id_data_array 
  51  from relax_errors import RelaxError, RelaxFileError, RelaxFileOverwriteError, RelaxInvalidSeqError, RelaxMissingBinaryError, RelaxNoInPathError, RelaxNonExecError 
  52  from relax_warnings import RelaxWarning, RelaxFileEmptyWarning 
  53   
  54   
  55   
56 -def delete(file_name, dir=None, fail=True):
57 """Deleting the given file, taking into account missing compression extensions. 58 59 @param file_name: The name of the file to delete. 60 @type file_name: str 61 @keyword dir: The directory containing the file. 62 @type dir: None or str 63 @keyword fail: A flag which if True will cause RelaxFileError to be raised. 64 @type fail: bool 65 @raises RelaxFileError: If the file does not exist, and fail is set to true. 66 """ 67 68 # File path. 69 file_path = get_file_path(file_name, dir) 70 71 # Test if the file exists and determine the compression type. 72 if access(file_path, F_OK): 73 pass 74 elif access(file_path + '.bz2', F_OK): 75 file_path = file_path + '.bz2' 76 elif access(file_path + '.gz', F_OK): 77 file_path = file_path + '.gz' 78 elif fail: 79 raise RelaxFileError(file_path) 80 else: 81 return 82 83 # Remove the file. 84 remove(file_path)
85 86
87 -def determine_compression(file_path):
88 """Function for determining the compression type, and for also testing if the file exists. 89 90 @param file_path: The full file path of the file. 91 @type file_path: str 92 @return: A tuple of the compression type and full path of the file (including its 93 extension). A value of 0 corresponds to no compression. Bzip2 compression 94 corresponds to a value of 1. Gzip compression corresponds to a value of 2. 95 @rtype: (int, str) 96 """ 97 98 # The file has been supplied without its compression extension. 99 if access(file_path, F_OK): 100 compress_type = 0 101 if search('.bz2$', file_path): 102 compress_type = 1 103 elif search('.gz$', file_path): 104 compress_type = 2 105 106 # The file has been supplied with the '.bz2' extension. 107 elif access(file_path + '.bz2', F_OK): 108 file_path = file_path + '.bz2' 109 compress_type = 1 110 111 # The file has been supplied with the '.gz' extension. 112 elif access(file_path + '.gz', F_OK): 113 file_path = file_path + '.gz' 114 compress_type = 2 115 116 # The file doesn't exist. 117 else: 118 raise RelaxFileError(file_path) 119 120 # Return the compression type. 121 return compress_type, file_path
122 123
124 -def extract_data(file=None, dir=None, file_data=None, sep=None):
125 """Return all data in the file as a list of lines where each line is a list of line elements. 126 127 @param file: The file to extract the data from. 128 @type file: str or file object 129 @param dir: The path where the file is located. If None and the file argument is a 130 string, then the current directory is assumed. 131 @type dir: str or None 132 @param file_data: If the file data has already been extracted from the file, it can be 133 passed into this function using this argument. If data is supplied 134 here, then the file_name and dir args are ignored. 135 @type file_data: list of str 136 @param sep: The character separating the columns in the file data. If None, then 137 whitespace is assumed. 138 @type sep: str 139 @return: The file data. 140 @rtype: list of lists of str 141 """ 142 143 # Data not already extracted from the file. 144 if not file_data: 145 # Open the file. 146 if isinstance(file, str): 147 file = open_read_file(file_name=file, dir=dir) 148 149 # Read lines. 150 file_data = file.readlines() 151 152 # Create a data structure from the contents of the file split by either whitespace or the separator, sep. 153 data = [] 154 for i in xrange(len(file_data)): 155 if sep: 156 row = split(file_data[i], sep) 157 else: 158 row = split(file_data[i]) 159 data.append(row) 160 return data 161 162 # Close the file. 163 if not file_data: 164 file.close()
165 166
167 -def file_root(file_path):
168 """Return the root file name, striped of path and extension details. 169 170 @param file_path: The full path to the file. 171 @type file_path: str 172 @return: The file root (with all directories and the extension stripped away). 173 @rtype: str 174 """ 175 176 root, ext = splitext(file_path) 177 return basename(root)
178 179
180 -def get_file_path(file_name=None, dir=None):
181 """Generate and expand the full file path. 182 183 @param file_name: The name of the file to extract the data from. 184 @type file_name: str 185 @param dir: The path where the file is located. If None, then the current directory is 186 assumed. 187 @type dir: str 188 @return: The full file path. 189 @rtype: str 190 """ 191 192 # File name. 193 file_path = file_name 194 195 # Add the directory. 196 if dir: 197 file_path = dir + sep + file_path 198 199 # Expand any ~ characters. 200 if file_path: # Catch a file path of None, as expanduser can't handle this. 201 file_path = expanduser(file_path) 202 203 # Return the file path. 204 return file_path
205 206
207 -def io_streams_restore(verbosity=1):
208 """Restore all IO streams to the Python defaults. 209 210 @param verbosity: The verbosity level. 211 @type verbosity: int 212 """ 213 214 # Print out. 215 if verbosity: 216 print("Restoring the sys.stdin IO stream to the Python STDIN IO stream.") 217 print("Restoring the sys.stdout IO stream to the Python STDOUT IO stream.") 218 print("Restoring the sys.stderr IO stream to the Python STDERR IO stream.") 219 220 # Restore streams. 221 sys.stdin = sys.__stdin__ 222 sys.stdout = sys.__stdout__ 223 sys.stderr = sys.__stderr__
224 225
226 -def io_streams_log(file_name=None, dir=None, verbosity=1):
227 """Turn on logging, sending both STDOUT and STDERR streams to a file. 228 229 @param file_name: The name of the file. 230 @type file_name: str 231 @param dir: The path where the file is located. If None, then the current directory is 232 assumed. 233 @type dir: str 234 @param verbosity: The verbosity level. 235 @type verbosity: int 236 """ 237 238 # Log file. 239 log_file, file_path = open_write_file(file_name=file_name, dir=dir, force=True, verbosity=verbosity, return_path=True) 240 241 # Logging IO streams. 242 log_stdin = stdin 243 log_stdout = None 244 log_stderr = SplitIO() 245 246 # Print out. 247 if verbosity: 248 print("Redirecting the sys.stdin IO stream to the python stdin IO stream.") 249 print(("Redirecting the sys.stdout IO stream to the log file '%s'." % file_path)) 250 print(("Redirecting the sys.stderr IO stream to both the python stderr IO stream and the log file '%s'." % file_path)) 251 252 # Set the logging IO streams. 253 log_stdout = log_file 254 log_stderr.split(stderr, log_file) 255 256 # IO stream redirection. 257 sys.stdin = log_stdin 258 sys.stdout = log_stdout 259 sys.stderr = log_stderr
260 261
262 -def io_streams_tee(file_name=None, dir=None, compress_type=0, verbosity=1):
263 """Turn on teeing to split both STDOUT and STDERR streams and sending second part to a file. 264 265 @param file_name: The name of the file. 266 @type file_name: str 267 @param dir: The path where the file is located. If None, then the current directory 268 is assumed. 269 @type dir: str 270 @param compress_type: The compression type. The integer values correspond to the compression 271 type: 0, no compression; 1, Bzip2 compression; 2, Gzip compression. 272 @type compress_type: int 273 @param verbosity: The verbosity level. 274 @type verbosity: int 275 """ 276 277 # Tee file. 278 tee_file, file_path = open_write_file(file_name=file_name, dir=dir, force=True, compress_type=compress_type, verbosity=verbosity, return_path=1) 279 280 # Tee IO streams. 281 tee_stdin = stdin 282 tee_stdout = SplitIO() 283 tee_stderr = SplitIO() 284 285 # Print out. 286 if verbosity: 287 print("Redirecting the sys.stdin IO stream to the python stdin IO stream.") 288 print(("Redirecting the sys.stdout IO stream to both the python stdout IO stream and the log file '%s'." % file_path)) 289 print(("Redirecting the sys.stderr IO stream to both the python stderr IO stream and the log file '%s'." % file_path)) 290 291 # Set the tee IO streams. 292 tee_stdout.split(stdout, tee_file) 293 tee_stderr.split(stderr, tee_file) 294 295 # IO stream redirection. 296 sys.stdin = tee_stdin 297 sys.stdout = tee_stdout 298 sys.stderr = tee_stderr
299 300
301 -def mkdir_nofail(dir=None, verbosity=1):
302 """Create the given directory, or exit without raising an error if the directory exists. 303 304 @param dir: The directory to create. 305 @type dir: str 306 @param verbosity: The verbosity level. 307 @type verbosity: int 308 """ 309 310 # No directory given. 311 if dir == None: 312 return 313 314 # Make the directory. 315 try: 316 makedirs(dir) 317 except OSError: 318 if verbosity: 319 print(("Directory ." + sep + dir + " already exists.\n"))
320 321
322 -def open_read_file(file_name=None, dir=None, verbosity=1):
323 """Open the file 'file' and return all the data. 324 325 @param file_name: The name of the file to extract the data from. 326 @type file_name: str 327 @param dir: The path where the file is located. If None, then the current directory is 328 assumed. 329 @type dir: str 330 @param verbosity: The verbosity level. 331 @type verbosity: int 332 @return: The open file object. 333 @rtype: file object 334 """ 335 336 # A file descriptor object. 337 if isinstance(file_name, file): 338 # Nothing to do here! 339 return file_name 340 341 # Invalid file name. 342 if not file_name and not isinstance(file_name, str): 343 raise RelaxError("The file name " + repr(file_name) + " " + repr(type(file_name)) + " is invalid and cannot be opened.") 344 345 # File path. 346 file_path = get_file_path(file_name, dir) 347 348 # Test if the file exists and determine the compression type. 349 compress_type, file_path = determine_compression(file_path) 350 351 # Open the file for reading. 352 try: 353 if verbosity: 354 print(("Opening the file " + repr(file_path) + " for reading.")) 355 if compress_type == 0: 356 file_obj = open(file_path, 'r') 357 elif compress_type == 1: 358 if dep_check.bz2_module: 359 file_obj = BZ2File(file_path, 'r') 360 else: 361 raise RelaxError("Cannot open the file " + repr(file_path) + ", try uncompressing first. " + dep_check.bz2_module_message + ".") 362 elif compress_type == 2: 363 file_obj = GzipFile(file_path, 'r') 364 except IOError, message: 365 raise RelaxError("Cannot open the file " + repr(file_path) + ". " + message.args[1] + ".") 366 367 # Return the opened file. 368 return file_obj
369 370
371 -def open_write_file(file_name=None, dir=None, force=False, compress_type=0, verbosity=1, return_path=False):
372 """Function for opening a file for writing and creating directories if necessary. 373 374 @param file_name: The name of the file to extract the data from. 375 @type file_name: str 376 @param dir: The path where the file is located. If None, then the current directory 377 is assumed. 378 @type dir: str 379 @param force: Boolean argument which if True causes the file to be overwritten if it 380 already exists. 381 @type force: bool 382 @param compress_type: The compression type. The integer values correspond to the compression 383 type: 0, no compression; 1, Bzip2 compression; 2, Gzip compression. 384 @type compress_type: int 385 @param verbosity: The verbosity level. 386 @type verbosity: int 387 @param return_path: If True, the function will return a tuple of the file object and the 388 full file path. 389 @type return_path: bool 390 @return: The open, writable file object and, if the return_path is True, then the 391 full file path is returned as well. 392 @rtype: writable file object (if return_path, then a tuple of the writable file 393 and the full file path) 394 """ 395 396 # A file descriptor object. 397 if isinstance(file_name, file): 398 # Nothing to do here! 399 return file_name 400 401 # Something pretending to be a file object. 402 if hasattr(file_name, 'write'): 403 # Nothing to do here! 404 return file_name 405 406 # The null device. 407 if search('devnull', file_name): 408 # Devnull could not be imported! 409 if not dep_check.devnull_import: 410 raise RelaxError(dep_check.devnull_import_message + ". To use devnull, please upgrade to Python >= 2.4.") 411 412 # Print out. 413 if verbosity: 414 print("Opening the null device file for writing.") 415 416 # Open the null device. 417 file_obj = open(devnull, 'w') 418 419 # Return the file. 420 if return_path: 421 return file_obj, None 422 else: 423 return file_obj 424 425 # Create the directories. 426 mkdir_nofail(dir, verbosity=0) 427 428 # File path. 429 file_path = get_file_path(file_name, dir) 430 431 # Bzip2 compression. 432 if compress_type == 1 and not search('.bz2$', file_path): 433 # Bz2 module exists. 434 if dep_check.bz2_module: 435 file_path = file_path + '.bz2' 436 437 # Switch to gzip compression. 438 else: 439 warn(RelaxWarning("Cannot use Bzip2 compression, using gzip compression instead. " + dep_check.bz2_module_message + ".")) 440 compress_type = 2 441 442 # Gzip compression. 443 if compress_type == 2 and not search('.gz$', file_path): 444 file_path = file_path + '.gz' 445 446 # Fail if the file already exists and the force flag is set to 0. 447 if access(file_path, F_OK) and not force: 448 raise RelaxFileOverwriteError(file_path, 'force flag') 449 450 # Open the file for writing. 451 try: 452 if verbosity: 453 print(("Opening the file " + repr(file_path) + " for writing.")) 454 if compress_type == 0: 455 file_obj = open(file_path, 'w') 456 elif compress_type == 1: 457 file_obj = BZ2File(file_path, 'w') 458 elif compress_type == 2: 459 file_obj = GzipFile(file_path, 'w') 460 except IOError, message: 461 raise RelaxError("Cannot open the file " + repr(file_path) + ". " + message.args[1] + ".") 462 463 # Return the opened file. 464 if return_path: 465 return file_obj, file_path 466 else: 467 return file_obj
468 469
470 -def read_spin_data(file=None, dir=None, file_data=None, spin_id_col=None, mol_name_col=None, res_num_col=None, res_name_col=None, spin_num_col=None, spin_name_col=None, data_col=None, error_col=None, sep=None, spin_id=None):
471 """Generator function for reading the spin specific data from file. 472 473 Description 474 =========== 475 476 This function reads a columnar formatted file where each line corresponds to a spin system. Spin identification is either through a spin ID string or through columns containing the molecule name, residue name and number, and/or spin name and number. 477 478 479 @keyword file: The name of the file to open. 480 @type file: str 481 @keyword dir: The directory containing the file (defaults to the current directory if None). 482 @type dir: str or None 483 @keyword file_data: An alternative to opening a file, if the data already exists in the correct format. The format is a list of lists where the first index corresponds to the row and the second the column. 484 @type file_data: list of lists 485 @keyword spin_id_col: The column containing the spin ID strings. If supplied, the mol_name_col, res_name_col, res_num_col, spin_name_col, and spin_num_col arguments must be none. 486 @type spin_id_col: int or None 487 @keyword mol_name_col: The column containing the molecule name information. If supplied, spin_id_col must be None. 488 @type mol_name_col: int or None 489 @keyword res_name_col: The column containing the residue name information. If supplied, spin_id_col must be None. 490 @type res_name_col: int or None 491 @keyword res_num_col: The column containing the residue number information. If supplied, spin_id_col must be None. 492 @type res_num_col: int or None 493 @keyword spin_name_col: The column containing the spin name information. If supplied, spin_id_col must be None. 494 @type spin_name_col: int or None 495 @keyword spin_num_col: The column containing the spin number information. If supplied, spin_id_col must be None. 496 @type spin_num_col: int or None 497 @keyword data_col: The column containing the data. 498 @type data_col: int or None 499 @keyword error_col: The column containing the errors. 500 @type error_col: int or None 501 @keyword sep: The column separator which, if None, defaults to whitespace. 502 @type sep: str or None 503 @keyword spin_id: The spin ID string used to restrict data loading to a subset of all spins. 504 @type spin_id: None or str 505 @return: A list of the spin specific data is yielded. The format is a list consisting of the spin ID string, the data value (if data_col is give), and the error value (if error_col is given). If both data_col and error_col are None, then the spin ID string is simply yielded. 506 @rtype: str, list of [str, float], or list of [str, float, float] 507 """ 508 509 # Argument tests. 510 col_args = [spin_id_col, mol_name_col, res_name_col, res_num_col, spin_name_col, spin_num_col, data_col, error_col] 511 col_arg_names = ['spin_id_col', 'mol_name_col', 'res_name_col', 'res_num_col', 'spin_name_col', 'spin_num_col', 'data_col', 'error_col'] 512 for i in range(len(col_args)): 513 if col_args[i] == 0: 514 raise RelaxError, "The '%s' argument cannot be zero, column numbering starts at one." % col_arg_names[i] 515 if spin_id_col and (mol_name_col or res_name_col or res_num_col or spin_name_col or spin_num_col): 516 raise RelaxError, "If the 'spin_id_col' argument has been supplied, then the mol_name_col, res_name_col, res_num_col, spin_name_col, and spin_num_col must all be set to None." 517 518 # Minimum number of columns. 519 min_col_num = max(spin_id_col, mol_name_col, res_num_col, res_name_col, spin_num_col, spin_name_col, data_col, error_col) 520 521 # Extract the data from the file. 522 if not file_data: 523 # Extract. 524 file_data = extract_data(file, dir) 525 526 # Strip the data of all comments and empty lines. 527 if spin_id_col != None: 528 file_data = strip(file_data, comments=False) 529 else: 530 file_data = strip(file_data) 531 532 # No data! 533 if not file_data: 534 warn(RelaxFileEmptyWarning(file)) 535 return 536 537 # Yield the data, spin by spin. 538 missing_data = True 539 for line in file_data: 540 # Validate the sequence. 541 try: 542 generic_fns.sequence.validate_sequence(line, spin_id_col=spin_id_col, mol_name_col=mol_name_col, res_num_col=res_num_col, res_name_col=res_name_col, spin_num_col=spin_num_col, spin_name_col=spin_name_col, data_col=data_col, error_col=error_col) 543 except RelaxInvalidSeqError, msg: 544 # Extract the message string, without the RelaxError bit. 545 string = msg.__str__()[12:-1] 546 547 # Give a warning. 548 warn(RelaxWarning(string)) 549 550 # Skip the line. 551 continue 552 553 # Generate the spin ID string. 554 if spin_id_col: 555 id = line[spin_id_col-1] 556 else: 557 id = generate_spin_id_data_array(data=line, mol_name_col=mol_name_col, res_num_col=res_num_col, res_name_col=res_name_col, spin_num_col=spin_num_col, spin_name_col=spin_name_col) 558 559 # Invalid spin ID. 560 if id == '#': 561 warn(RelaxWarning("Invalid spin ID, skipping the line %s" % line)) 562 continue 563 564 # Convert the data. 565 value = None 566 if data_col: 567 try: 568 # None. 569 if line[data_col-1] == 'None': 570 value = None 571 572 # A float. 573 else: 574 value = float(line[data_col-1]) 575 576 # Bad data. 577 except ValueError: 578 warn(RelaxWarning("Invalid data, skipping the line %s" % line)) 579 continue 580 581 # Convert the errors. 582 error = None 583 if error_col: 584 try: 585 # None. 586 if line[error_col-1] == 'None': 587 error = None 588 589 # A float. 590 else: 591 error = float(line[error_col-1]) 592 593 # Bad data. 594 except ValueError: 595 warn(RelaxWarning("Invalid errors, skipping the line %s" % line)) 596 continue 597 598 # Right, data is OK and exists. 599 missing_data = False 600 601 # Yield the data. 602 if data_col and error_col: 603 yield [id, value, error] 604 elif data_col: 605 yield [id, value] 606 elif error_col: 607 yield [id, error] 608 else: 609 yield id 610 611 # Hmmm, no data! 612 if missing_data: 613 raise RelaxError("No corresponding data could be found within the file.")
614 615
616 -def strip(data, comments=True):
617 """Remove all comment and empty lines from the file data structure. 618 619 @param data: The file data to clean up. 620 @type data: list of lists of str 621 @keyword comments: A flag which if True will cause comments to be deleted. 622 @type comments: bool 623 @return: The input data with the empty and comment lines removed. 624 @rtype: list of lists of str 625 """ 626 627 # Initialise the new data array. 628 new = [] 629 630 # Loop over the data. 631 for i in xrange(len(data)): 632 # Empty lines. 633 if len(data[i]) == 0: 634 continue 635 636 # Comment lines. 637 if comments and search("^#", data[i][0]): 638 continue 639 640 # Data lines. 641 new.append(data[i]) 642 643 # Return the new data structure. 644 return new
645 646
647 -def test_binary(binary):
648 """Function for testing that the binary string corresponds to a valid executable file. 649 650 @param binary: The name or path of the binary executable file. 651 @type binary: str 652 """ 653 654 # Path separator RE string. 655 if altsep: 656 path_sep = '[' + sep + altsep + ']' 657 else: 658 path_sep = sep 659 660 # The full path of the program has been given (if a directory separatory has been supplied). 661 if search(path_sep, binary): 662 # Test that the binary exists. 663 if not access(binary, F_OK): 664 raise RelaxMissingBinaryError(binary) 665 666 # Test that if the binary is executable. 667 if not access(binary, X_OK): 668 raise RelaxNonExecError(binary) 669 670 # The path to the binary has not been given. 671 else: 672 # Get the PATH environmental variable. 673 path = getenv('PATH') 674 675 # Split PATH by the path separator. 676 path_list = split(path, pathsep) 677 678 # Test that the binary exists within the system path (and exit this function instantly once it has been found). 679 for path in path_list: 680 if access(path + sep + binary, F_OK): 681 return 682 683 # The binary is not located in the system path! 684 raise RelaxNoInPathError(binary)
685 686
687 -def write_spin_data(file, dir=None, sep=None, spin_ids=None, mol_names=None, res_nums=None, res_names=None, spin_nums=None, spin_names=None, force=False, data=None, data_name=None, data_length=20, data_format=None, error=None, error_name=None, error_length=20, error_format=None):
688 """Generator function for reading the spin specific data from file. 689 690 Description 691 =========== 692 693 This function writes a columnar formatted file where each line corresponds to a spin system. 694 Spin identification is either through a spin ID string or through columns containing the 695 molecule name, residue name and number, and/or spin name and number. 696 697 698 @param file: The name of the file to write the data to (or alternatively an already opened file object). 699 @type file: str or file object 700 @keyword dir: The directory to place the file into (defaults to the current directory if None and the file argument is not a file object). 701 @type dir: str or None 702 @keyword sep: The column separator which, if None, defaults to whitespace. 703 @type sep: str or None 704 @keyword spin_ids: The list of spin ID strings. 705 @type spin_ids: None or list of str 706 @keyword mol_names: The list of molecule names. 707 @type mol_names: None or list of str 708 @keyword res_nums: The list of residue numbers. 709 @type res_nums: None or list of int 710 @keyword res_names: The list of residue names. 711 @type res_names: None or list of str 712 @keyword spin_nums: The list of spin numbers. 713 @type spin_nums: None or list of int 714 @keyword spin_names: The list of spin names. 715 @type spin_names: None or list of str 716 @keyword force: A flag which if True will cause an existing file to be overwritten. 717 @type force: bool 718 @keyword data: A list of the data to write out. The first dimension corresponds to the spins. A second dimension can also be given if multiple data sets across multiple columns are desired. 719 @type data: list or list of lists 720 @keyword data_name: A name corresponding to the data argument. If the data argument is a list of lists, then this must also be a list with the same length as the second dimension of the data arg. 721 @type data_name: str or list of str 722 @keyword data_length: The length of the data columns. 723 @type data_length: int 724 @keyword data_format: The optional python formatting string for the data columns, e.g. "%-30s". 725 @type data_format: None or str 726 @keyword error: A list of the errors to write out. The first dimension corresponds to the spins. A second dimension can also be given if multiple data sets across multiple columns are desired. These will be inter-dispersed between the data columns, if the data is given. If the data arg is not None, then this must have the same dimensions as that object. 727 @type error: list or list of lists 728 @keyword error_name: A name corresponding to the error argument. If the error argument is a list of lists, then this must also be a list with the same length at the second dimension of the error arg. 729 @type error_name: str or list of str 730 @keyword error_length: The length of the error columns. 731 @type error_length: int 732 @keyword error_format: The optional python formatting string for the error columns, e.g. "%-30s". 733 @type error_format: None or str 734 """ 735 736 # Data argument tests. 737 if data: 738 # Data is a list of lists. 739 if isinstance(data[0], list): 740 # Data and data_name don't match. 741 if not isinstance(data_name, list): 742 raise RelaxError("The data_name arg '%s' must be a list as the data argument is a list of lists." % data_name) 743 744 # Error doesn't match. 745 if error and (len(data) != len(error) or len(data[0]) != len(error[0])): 746 raise RelaxError("The data arg:\n%s\n\ndoes not have the same dimensions as the error arg:\n%s." % (data, error)) 747 748 # Data is a simple list. 749 else: 750 # Data and data_name don't match. 751 if not isinstance(data_name, str): 752 raise RelaxError("The data_name arg '%s' must be a string as the data argument is a simple list." % data_name) 753 754 # Error doesn't match. 755 if error and len(data) != len(error): 756 raise RelaxError("The data arg:\n%s\n\ndoes not have the same dimensions as the error arg:\n%s." % (data, error)) 757 758 # Error argument tests. 759 if error: 760 # Error is a list of lists. 761 if isinstance(error[0], list): 762 # Error and error_name don't match. 763 if not isinstance(error_name, list): 764 raise RelaxError("The error_name arg '%s' must be a list as the error argument is a list of lists." % error_name) 765 766 # Error is a simple list. 767 else: 768 # Error and error_name don't match. 769 if not isinstance(error_name, str): 770 raise RelaxError("The error_name arg '%s' must be a string as the error argument is a simple list." % error_name) 771 772 # Number of spins check. 773 args = [spin_ids, mol_names, res_nums, res_names, spin_nums, spin_names] 774 arg_names = ['spin_ids', 'mol_names', 'res_nums', 'res_names', 'spin_nums', 'spin_names'] 775 N = None 776 first_arg = None 777 first_arg_name = None 778 for i in range(len(args)): 779 if isinstance(args[i], list): 780 # First list match. 781 if N == None: 782 N = len(args[i]) 783 first_arg = args[i] 784 first_arg_name = arg_names[i] 785 786 # Length check. 787 if len(args[i]) != N: 788 raise RelaxError("The %s and %s arguments do not have the same number of spins ('%s' vs. '%s' respectively)." % (first_arg_name, arg_names[i], len(first_arg), len(args[i]))) 789 790 # Nothing?!? 791 if N == None: 792 raise RelaxError("No spin ID data is present.") 793 794 # Data and error length check. 795 if data and len(data) != N: 796 raise RelaxError("The %s and data arguments do not have the same number of spins ('%s' vs. '%s' respectively)." % (first_arg_name, len(first_arg), len(data))) 797 if error and len(error) != N: 798 raise RelaxError("The %s and error arguments do not have the same number of spins ('%s' vs. '%s' respectively)." % (first_arg_name, len(first_arg), len(error))) 799 800 # The spin arguments. 801 args = [spin_ids, mol_names, res_nums, res_names, spin_nums, spin_names] 802 arg_names = ['spin_id', 'mol_name', 'res_num', 'res_name', 'spin_num', 'spin_name'] 803 804 # No special separator character. 805 if sep == None: 806 sep = '' 807 808 # Open the file. 809 file = open_write_file(file_name=file, dir=dir, force=force) 810 811 # The spin ID column lengths. 812 len_args = [10] * 6 813 for i in range(len(args)): 814 # Different separator. 815 if sep: 816 len_args[i] = 0 817 continue 818 819 # Minimum width of the header name lengths. 820 if args[i] and len(arg_names[i]) > len_args[i]: 821 len_args[i] = len(arg_names[i]) + 2 822 823 # Minimum width of the spin ID data. 824 for spin_index in range(N): 825 if args[i] and len(repr(args[i][spin_index])) > len_args[i]: 826 len_args[i] = len(repr(args[i][spin_index])) + 1 827 828 # Data and error formatting strings. 829 if sep: 830 data_length = '' 831 data_head_format = "%%-%ss" % data_length 832 if not data_format: 833 data_format = "%%%ss" % data_length 834 error_head_format = "%%-%ss" % error_length 835 if not error_format: 836 error_format = "%%%ss" % error_length 837 838 839 # Header. 840 ######### 841 842 # Init. 843 file.write("\n") 844 prefix = '# ' 845 shift = -2 846 847 # The spin ID info. 848 for i in range(len(args)): 849 if args[i]: 850 file.write(("%s%%-%ss%s" % (prefix, len_args[i]+shift, sep)) % arg_names[i]) 851 prefix = ' ' 852 shift = 0 853 854 # The data. 855 if data: 856 # List of lists. 857 if isinstance(data[0], list): 858 # Loop over the list. 859 for i in range(len(data[0])): 860 # The data. 861 file.write((prefix+data_head_format+sep) % data_name[i]) 862 863 # The error. 864 if error: 865 file.write((prefix+error_head_format+sep) % error_name[i]) 866 867 # Simple list. 868 else: 869 # The data. 870 file.write((prefix+data_head_format+sep) % data_name) 871 872 # The error. 873 if error: 874 file.write((prefix+error_head_format+sep) % error_name) 875 876 # Only errors. 877 elif error: 878 # List of lists. 879 if isinstance(error[0], list): 880 for i in range(len(error[0])): 881 file.write((prefix+error_head_format+sep) % error_name[i]) 882 883 # Simple list. 884 else: 885 file.write((prefix+error_head_format+sep) % error_name) 886 887 # Terminate the line. 888 file.write("\n") 889 890 891 # Spin specific data. 892 ##################### 893 894 # Loop over the spins. 895 for spin_index in range(N): 896 # The prefix. 897 prefix = '' 898 899 # The spin ID info. 900 for i in range(len(args)): 901 if args[i]: 902 file.write(("%s%%-%ss%s" % (prefix, len_args[i], sep)) % args[i][spin_index]) 903 prefix = ' ' 904 905 # The data. 906 if data: 907 # List of lists. 908 if isinstance(data[0], list): 909 # Loop over the list. 910 for i in range(len(data[0])): 911 # The data. 912 file.write((prefix+data_format+sep) % data[spin_index][i]) 913 914 # The error. 915 if error: 916 file.write((prefix+error_format+sep) % error[spin_index][i]) 917 918 # Simple list. 919 else: 920 # The data. 921 file.write((prefix+data_format+sep) % data[spin_index]) 922 923 # The error. 924 if error: 925 file.write((prefix+error_format+sep) % error[spin_index]) 926 927 # Only errors. 928 elif error: 929 # List of lists. 930 if isinstance(error[0], list): 931 for i in range(len(error[0])): 932 file.write((prefix+error_format+sep) % error[spin_index][i]) 933 934 # Simple list. 935 else: 936 file.write((prefix+error_format+sep) % error[spin_index]) 937 938 # End of the line. 939 file.write("\n")
940 941 942
943 -class DummyFileObject:
944 - def __init__(self):
945 """Set up the dummy object to act as a file object.""" 946 947 # Initialise an object for adding the string from all write calls to. 948 self.data = '' 949 950 # Set the closed flag. 951 self.closed = False
952 953
954 - def close(self):
955 """A method for 'closing' this object.""" 956 957 # Set the closed flag. 958 self.closed = True
959 960
961 - def write(self, str):
962 """Mimic the file object write() method so that this class can be used as a file object. 963 964 @param str: The string to be written. 965 @type str: str 966 """ 967 968 # Check if the file is closed. 969 if self.closed: 970 raise ValueError('I/O operation on closed file') 971 972 # Append the string to the data object. 973 self.data = self.data + str
974 975
976 - def readlines(self):
977 """Mimic the file object readlines() method. 978 979 This method works even if this dummy file object is closed! 980 981 982 @return: The contents of the file object separated by newline characters. 983 @rtype: list of str 984 """ 985 986 # Split up the string. 987 lines = split(self.data, '\n') 988 989 # Remove the last line if empty. 990 if lines[-1] == '': 991 lines.pop() 992 993 # Loop over the lines, re-adding the newline character to match the file object readlines() method. 994 for i in xrange(len(lines)): 995 lines[i] = lines[i] + '\n' 996 997 # Return the file lines. 998 return lines
999 1000 1001
1002 -class SplitIO:
1003 - def __init__(self):
1004 """Class for splitting an IO stream to two outputs."""
1005 1006
1007 - def flush(self):
1008 """Flush all streams.""" 1009 1010 # Call the streams' methods. 1011 self.stream1.flush() 1012 self.stream2.flush()
1013 1014
1015 - def isatty(self):
1016 """Check that both streams are TTYs. 1017 1018 @return: True, only if both streams are TTYs. 1019 @rtype: bool 1020 """ 1021 1022 # Check both streams. 1023 return self.stream1.isatty() & self.stream2.isatty()
1024 1025
1026 - def split(self, stream1, stream2):
1027 """Function for setting the streams.""" 1028 1029 # Arguments. 1030 self.stream1 = stream1 1031 self.stream2 = stream2
1032 1033
1034 - def write(self, text):
1035 """Replacement write function.""" 1036 1037 # Write to stream1. 1038 self.stream1.write(text) 1039 1040 # Write to stream2. 1041 self.stream2.write(text)
1042