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 relax_errors import RelaxError, RelaxFileError, RelaxFileOverwriteError, RelaxInvalidSeqError, RelaxMissingBinaryError, RelaxNoInPathError, RelaxNonExecError 51 from relax_warnings import RelaxWarning, RelaxFileEmptyWarning 52 53 54

55 -def delete(file_name, dir=None, fail=True):

56 """Deleting the given file, taking into account missing compression extensions. 57 58 @param file_name: The name of the file to delete. 59 @type file_name: str 60 @keyword dir: The directory containing the file. 61 @type dir: None or str 62 @keyword fail: A flag which if True will cause RelaxFileError to be raised. 63 @type fail: bool 64 @raises RelaxFileError: If the file does not exist, and fail is set to true. 65 """ 66 67 # File path. 68 file_path = get_file_path(file_name, dir) 69 70 # Test if the file exists and determine the compression type. 71 if access(file_path, F_OK): 72 pass 73 elif access(file_path + '.bz2', F_OK): 74 file_path = file_path + '.bz2' 75 elif access(file_path + '.gz', F_OK): 76 file_path = file_path + '.gz' 77 elif fail: 78 raise RelaxFileError(file_path) 79 else: 80 return 81 82 # Remove the file. 83 remove(file_path)

84 85

86 -def determine_compression(file_path):

87 """Function for determining the compression type, and for also testing if the file exists. 88 89 @param file_path: The full file path of the file. 90 @type file_path: str 91 @return: A tuple of the compression type and full path of the file (including its 92 extension). A value of 0 corresponds to no compression. Bzip2 compression 93 corresponds to a value of 1. Gzip compression corresponds to a value of 2. 94 @rtype: (int, str) 95 """ 96 97 # The file has been supplied without its compression extension. 98 if access(file_path, F_OK): 99 compress_type = 0 100 if search('.bz2$', file_path): 101 compress_type = 1 102 elif search('.gz$', file_path): 103 compress_type = 2 104 105 # The file has been supplied with the '.bz2' extension. 106 elif access(file_path + '.bz2', F_OK): 107 file_path = file_path + '.bz2' 108 compress_type = 1 109 110 # The file has been supplied with the '.gz' extension. 111 elif access(file_path + '.gz', F_OK): 112 file_path = file_path + '.gz' 113 compress_type = 2 114 115 # The file doesn't exist. 116 else: 117 raise RelaxFileError(file_path) 118 119 # Return the compression type. 120 return compress_type, file_path

121 122

123 -def extract_data(file=None, dir=None, file_data=None, sep=None):

124 """Return all data in the file as a list of lines where each line is a list of line elements. 125 126 @param file: The file to extract the data from. 127 @type file: str or file object 128 @param dir: The path where the file is located. If None and the file argument is a 129 string, then the current directory is assumed. 130 @type dir: str or None 131 @param file_data: If the file data has already been extracted from the file, it can be 132 passed into this function using this argument. If data is supplied 133 here, then the file_name and dir args are ignored. 134 @type file_data: list of str 135 @param sep: The character separating the columns in the file data. If None, then 136 whitespace is assumed. 137 @type sep: str 138 @return: The file data. 139 @rtype: list of lists of str 140 """ 141 142 # Data not already extracted from the file. 143 if not file_data: 144 # Open the file. 145 if isinstance(file, str): 146 file = open_read_file(file_name=file, dir=dir) 147 148 # Read lines. 149 file_data = file.readlines() 150 151 # Create a data structure from the contents of the file split by either whitespace or the separator, sep. 152 data = [] 153 for i in xrange(len(file_data)): 154 if sep: 155 row = split(file_data[i], sep) 156 else: 157 row = split(file_data[i]) 158 data.append(row) 159 return data 160 161 # Close the file. 162 if not file_data: 163 file.close()

164 165

166 -def file_root(file_path):

167 """Return the root file name, striped of path and extension details. 168 169 @param file_path: The full path to the file. 170 @type file_path: str 171 @return: The file root (with all directories and the extension stripped away). 172 @rtype: str 173 """ 174 175 root, ext = splitext(file_path) 176 return basename(root)

177 178

179 -def get_file_path(file_name=None, dir=None):

180 """Generate and expand the full file path. 181 182 @param file_name: The name of the file to extract the data from. 183 @type file_name: str 184 @param dir: The path where the file is located. If None, then the current directory is 185 assumed. 186 @type dir: str 187 @return: The full file path. 188 @rtype: str 189 """ 190 191 # File name. 192 file_path = file_name 193 194 # Add the directory. 195 if dir: 196 file_path = dir + sep + file_path 197 198 # Expand any ~ characters. 199 if file_path: # Catch a file path of None, as expanduser can't handle this. 200 file_path = expanduser(file_path) 201 202 # Return the file path. 203 return file_path

204 205

206 -def io_streams_restore(verbosity=1):

207 """Restore all IO streams to the Python defaults. 208 209 @param verbosity: The verbosity level. 210 @type verbosity: int 211 """ 212 213 # Print out. 214 if verbosity: 215 print("Restoring the sys.stdin IO stream to the Python STDIN IO stream.") 216 print("Restoring the sys.stdout IO stream to the Python STDOUT IO stream.") 217 print("Restoring the sys.stderr IO stream to the Python STDERR IO stream.") 218 219 # Restore streams. 220 sys.stdin = sys.__stdin__ 221 sys.stdout = sys.__stdout__ 222 sys.stderr = sys.__stderr__

223 224

225 -def io_streams_log(file_name=None, dir=None, verbosity=1):

226 """Turn on logging, sending both STDOUT and STDERR streams to a file. 227 228 @param file_name: The name of the file. 229 @type file_name: str 230 @param dir: The path where the file is located. If None, then the current directory is 231 assumed. 232 @type dir: str 233 @param verbosity: The verbosity level. 234 @type verbosity: int 235 """ 236 237 # Log file. 238 log_file, file_path = open_write_file(file_name=file_name, dir=dir, force=True, verbosity=verbosity, return_path=True) 239 240 # Logging IO streams. 241 log_stdin = stdin 242 log_stdout = None 243 log_stderr = SplitIO() 244 245 # Print out. 246 if verbosity: 247 print("Redirecting the sys.stdin IO stream to the python stdin IO stream.") 248 print(("Redirecting the sys.stdout IO stream to the log file '%s'." % file_path)) 249 print(("Redirecting the sys.stderr IO stream to both the python stderr IO stream and the log file '%s'." % file_path)) 250 251 # Set the logging IO streams. 252 log_stdout = log_file 253 log_stderr.split(stderr, log_file) 254 255 # IO stream redirection. 256 sys.stdin = log_stdin 257 sys.stdout = log_stdout 258 sys.stderr = log_stderr

259 260

261 -def io_streams_tee(file_name=None, dir=None, compress_type=0, verbosity=1):

262 """Turn on teeing to split both STDOUT and STDERR streams and sending second part to a file. 263 264 @param file_name: The name of the file. 265 @type file_name: str 266 @param dir: The path where the file is located. If None, then the current directory 267 is assumed. 268 @type dir: str 269 @param compress_type: The compression type. The integer values correspond to the compression 270 type: 0, no compression; 1, Bzip2 compression; 2, Gzip compression. 271 @type compress_type: int 272 @param verbosity: The verbosity level. 273 @type verbosity: int 274 """ 275 276 # Tee file. 277 tee_file, file_path = open_write_file(file_name=file_name, dir=dir, force=True, compress_type=compress_type, verbosity=verbosity, return_path=1) 278 279 # Tee IO streams. 280 tee_stdin = stdin 281 tee_stdout = SplitIO() 282 tee_stderr = SplitIO() 283 284 # Print out. 285 if verbosity: 286 print("Redirecting the sys.stdin IO stream to the python stdin IO stream.") 287 print(("Redirecting the sys.stdout IO stream to both the python stdout IO stream and the log file '%s'." % file_path)) 288 print(("Redirecting the sys.stderr IO stream to both the python stderr IO stream and the log file '%s'." % file_path)) 289 290 # Set the tee IO streams. 291 tee_stdout.split(stdout, tee_file) 292 tee_stderr.split(stderr, tee_file) 293 294 # IO stream redirection. 295 sys.stdin = tee_stdin 296 sys.stdout = tee_stdout 297 sys.stderr = tee_stderr

298 299

300 -def mkdir_nofail(dir=None, verbosity=1):

301 """Create the given directory, or exit without raising an error if the directory exists. 302 303 @param dir: The directory to create. 304 @type dir: str 305 @param verbosity: The verbosity level. 306 @type verbosity: int 307 """ 308 309 # No directory given. 310 if dir == None: 311 return 312 313 # Make the directory. 314 try: 315 makedirs(dir) 316 except OSError: 317 if verbosity: 318 print(("Directory ." + sep + dir + " already exists.\n"))

319 320

321 -def open_read_file(file_name=None, dir=None, verbosity=1):

322 """Open the file 'file' and return all the data. 323 324 @param file_name: The name of the file to extract the data from. 325 @type file_name: str 326 @param dir: The path where the file is located. If None, then the current directory is 327 assumed. 328 @type dir: str 329 @param verbosity: The verbosity level. 330 @type verbosity: int 331 @return: The open file object. 332 @rtype: file object 333 """ 334 335 # A file descriptor object. 336 if isinstance(file_name, file): 337 # Nothing to do here! 338 return file_name 339 340 # Invalid file name. 341 if not file_name and not isinstance(file_name, str): 342 raise RelaxError("The file name " + repr(file_name) + " " + repr(type(file_name)) + " is invalid and cannot be opened.") 343 344 # File path. 345 file_path = get_file_path(file_name, dir) 346 347 # Test if the file exists and determine the compression type. 348 compress_type, file_path = determine_compression(file_path) 349 350 # Open the file for reading. 351 try: 352 if verbosity: 353 print(("Opening the file " + repr(file_path) + " for reading.")) 354 if compress_type == 0: 355 file_obj = open(file_path, 'r') 356 elif compress_type == 1: 357 if dep_check.bz2_module: 358 file_obj = BZ2File(file_path, 'r') 359 else: 360 raise RelaxError("Cannot open the file " + repr(file_path) + ", try uncompressing first. " + dep_check.bz2_module_message + ".") 361 elif compress_type == 2: 362 file_obj = GzipFile(file_path, 'r') 363 except IOError, message: 364 raise RelaxError("Cannot open the file " + repr(file_path) + ". " + message.args[1] + ".") 365 366 # Return the opened file. 367 return file_obj

368 369

370 -def open_write_file(file_name=None, dir=None, force=False, compress_type=0, verbosity=1, return_path=False):

371 """Function for opening a file for writing and creating directories if necessary. 372 373 @param file_name: The name of the file to extract the data from. 374 @type file_name: str 375 @param dir: The path where the file is located. If None, then the current directory 376 is assumed. 377 @type dir: str 378 @param force: Boolean argument which if True causes the file to be overwritten if it 379 already exists. 380 @type force: bool 381 @param compress_type: The compression type. The integer values correspond to the compression 382 type: 0, no compression; 1, Bzip2 compression; 2, Gzip compression. 383 @type compress_type: int 384 @param verbosity: The verbosity level. 385 @type verbosity: int 386 @param return_path: If True, the function will return a tuple of the file object and the 387 full file path. 388 @type return_path: bool 389 @return: The open, writable file object and, if the return_path is True, then the 390 full file path is returned as well. 391 @rtype: writable file object (if return_path, then a tuple of the writable file 392 and the full file path) 393 """ 394 395 # A file descriptor object. 396 if isinstance(file_name, file): 397 # Nothing to do here! 398 return file_name 399 400 # Something pretending to be a file object. 401 if hasattr(file_name, 'write'): 402 # Nothing to do here! 403 return file_name 404 405 # The null device. 406 if search('devnull', file_name): 407 # Devnull could not be imported! 408 if not dep_check.devnull_import: 409 raise RelaxError(dep_check.devnull_import_message + ". To use devnull, please upgrade to Python >= 2.4.") 410 411 # Print out. 412 if verbosity: 413 print("Opening the null device file for writing.") 414 415 # Open the null device. 416 file_obj = open(devnull, 'w') 417 418 # Return the file. 419 if return_path: 420 return file_obj, None 421 else: 422 return file_obj 423 424 # Create the directories. 425 mkdir_nofail(dir, verbosity=0) 426 427 # File path. 428 file_path = get_file_path(file_name, dir) 429 430 # Bzip2 compression. 431 if compress_type == 1 and not search('.bz2$', file_path): 432 # Bz2 module exists. 433 if dep_check.bz2_module: 434 file_path = file_path + '.bz2' 435 436 # Switch to gzip compression. 437 else: 438 warn(RelaxWarning("Cannot use Bzip2 compression, using gzip compression instead. " + dep_check.bz2_module_message + ".")) 439 compress_type = 2 440 441 # Gzip compression. 442 if compress_type == 2 and not search('.gz$', file_path): 443 file_path = file_path + '.gz' 444 445 # Fail if the file already exists and the force flag is set to 0. 446 if access(file_path, F_OK) and not force: 447 raise RelaxFileOverwriteError(file_path, 'force flag') 448 449 # Open the file for writing. 450 try: 451 if verbosity: 452 print(("Opening the file " + repr(file_path) + " for writing.")) 453 if compress_type == 0: 454 file_obj = open(file_path, 'w') 455 elif compress_type == 1: 456 file_obj = BZ2File(file_path, 'w') 457 elif compress_type == 2: 458 file_obj = GzipFile(file_path, 'w') 459 except IOError, message: 460 raise RelaxError("Cannot open the file " + repr(file_path) + ". " + message.args[1] + ".") 461 462 # Return the opened file. 463 if return_path: 464 return file_obj, file_path 465 else: 466 return file_obj

467 468

469 -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):

470 """Generator function for reading the spin specific data from file. 471 472 Description 473 =========== 474 475 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. 476 477 478 @keyword file: The name of the file to open. 479 @type file: str 480 @keyword dir: The directory containing the file (defaults to the current directory if None). 481 @type dir: str or None 482 @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. 483 @type file_data: list of lists 484 @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. 485 @type spin_id_col: int or None 486 @keyword mol_name_col: The column containing the molecule name information. If supplied, spin_id_col must be None. 487 @type mol_name_col: int or None 488 @keyword res_name_col: The column containing the residue name information. If supplied, spin_id_col must be None. 489 @type res_name_col: int or None 490 @keyword res_num_col: The column containing the residue number information. If supplied, spin_id_col must be None. 491 @type res_num_col: int or None 492 @keyword spin_name_col: The column containing the spin name information. If supplied, spin_id_col must be None. 493 @type spin_name_col: int or None 494 @keyword spin_num_col: The column containing the spin number information. If supplied, spin_id_col must be None. 495 @type spin_num_col: int or None 496 @keyword data_col: The column containing the data. 497 @type data_col: int or None 498 @keyword error_col: The column containing the errors. 499 @type error_col: int or None 500 @keyword sep: The column separator which, if None, defaults to whitespace. 501 @type sep: str or None 502 @keyword spin_id: The spin ID string used to restrict data loading to a subset of all spins. 503 @type spin_id: None or str 504 @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. 505 @rtype: str, list of [str, float], or list of [str, float, float] 506 """ 507 508 # Argument tests. 509 col_args = [spin_id_col, mol_name_col, res_name_col, res_num_col, spin_name_col, spin_num_col, data_col, error_col] 510 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'] 511 for i in range(len(col_args)): 512 if col_args[i] == 0: 513 raise RelaxError, "The '%s' argument cannot be zero, column numbering starts at one." % col_arg_names[i] 514 if spin_id_col and (mol_name_col or res_name_col or res_num_col or spin_name_col or spin_num_col): 515 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." 516 517 # Minimum number of columns. 518 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) 519 520 # Extract the data from the file. 521 if not file_data: 522 # Extract. 523 file_data = extract_data(file, dir) 524 525 # Strip the data of all comments and empty lines. 526 if spin_id_col != None: 527 file_data = strip(file_data, comments=False) 528 else: 529 file_data = strip(file_data) 530 531 # No data! 532 if not file_data: 533 warn(RelaxFileEmptyWarning(file)) 534 return 535 536 # Yield the data, spin by spin. 537 missing_data = True 538 for line in file_data: 539 # Validate the sequence. 540 try: 541 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) 542 except RelaxInvalidSeqError, msg: 543 # Extract the message string, without the RelaxError bit. 544 string = msg.__str__()[12:-1] 545 546 # Give a warning. 547 warn(RelaxWarning(string)) 548 549 # Skip the line. 550 continue 551 552 # Get the spin data from the ID. 553 if spin_id_col: 554 # Invalid spin ID. 555 if line[spin_id_col-1] == '#': 556 warn(RelaxWarning("Invalid spin ID, skipping the line %s" % line)) 557 continue 558 559 mol_name, res_num, res_name, spin_num, spin_name = generic_fns.mol_res_spin.spin_id_to_data_list(line[spin_id_col-1]) 560 561 # Convert the spin data. 562 else: 563 # The molecule. 564 mol_name = None 565 if mol_name_col != None and line[mol_name_col-1] != 'None': 566 mol_name = line[mol_name_col-1] 567 568 # The residue number, catching bad values. 569 res_num = None 570 if res_num_col != None: 571 try: 572 if line[res_num_col-1] == 'None': 573 res_num = None 574 else: 575 res_num = int(line[res_num_col-1]) 576 except ValueError: 577 warn(RelaxWarning("Invalid residue number, skipping the line %s" % line)) 578 continue 579 580 # The residue name. 581 res_name = None 582 if res_name_col != None and line[res_name_col-1] != 'None': 583 res_name = line[res_name_col-1] 584 585 # The spin number, catching bad values. 586 spin_num = None 587 if spin_num_col != None: 588 try: 589 if line[spin_num_col-1] == 'None': 590 spin_num = None 591 else: 592 spin_num = int(line[spin_num_col-1]) 593 except ValueError: 594 warn(RelaxWarning("Invalid spin number, skipping the line %s" % line)) 595 continue 596 597 # The spin name. 598 spin_name = None 599 if spin_name_col != None and line[spin_name_col-1] != 'None': 600 spin_name = line[spin_name_col-1] 601 602 # Convert the data. 603 value = None 604 if data_col != None: 605 try: 606 # None. 607 if line[data_col-1] == 'None': 608 value = None 609 610 # A float. 611 else: 612 value = float(line[data_col-1]) 613 614 # Bad data. 615 except ValueError: 616 warn(RelaxWarning("Invalid data, skipping the line %s" % line)) 617 continue 618 619 # Convert the errors. 620 error = None 621 if error_col != None: 622 try: 623 # None. 624 if line[error_col-1] == 'None': 625 error = None 626 627 # A float. 628 else: 629 error = float(line[error_col-1]) 630 631 # Bad data. 632 except ValueError: 633 warn(RelaxWarning("Invalid errors, skipping the line %s" % line)) 634 continue 635 636 # Right, data is OK and exists. 637 missing_data = False 638 639 # Yield the data. 640 if data_col and error_col: 641 yield mol_name, res_num, res_name, spin_num, spin_name, value, error 642 elif data_col: 643 yield mol_name, res_num, res_name, spin_num, spin_name, value 644 elif error_col: 645 yield mol_name, res_num, res_name, spin_num, spin_name, error 646 else: 647 yield mol_name, res_num, res_name, spin_num, spin_name 648 649 # Hmmm, no data! 650 if missing_data: 651 raise RelaxError("No corresponding data could be found within the file.")

652 653

654 -def strip(data, comments=True):

655 """Remove all comment and empty lines from the file data structure. 656 657 @param data: The file data to clean up. 658 @type data: list of lists of str 659 @keyword comments: A flag which if True will cause comments to be deleted. 660 @type comments: bool 661 @return: The input data with the empty and comment lines removed. 662 @rtype: list of lists of str 663 """ 664 665 # Initialise the new data array. 666 new = [] 667 668 # Loop over the data. 669 for i in xrange(len(data)): 670 # Empty lines. 671 if len(data[i]) == 0: 672 continue 673 674 # Comment lines. 675 if comments and search("^#", data[i][0]): 676 continue 677 678 # Data lines. 679 new.append(data[i]) 680 681 # Return the new data structure. 682 return new

683 684

685 -def test_binary(binary):

686 """Function for testing that the binary string corresponds to a valid executable file. 687 688 @param binary: The name or path of the binary executable file. 689 @type binary: str 690 """ 691 692 # Path separator RE string. 693 if altsep: 694 path_sep = '[' + sep + altsep + ']' 695 else: 696 path_sep = sep 697 698 # The full path of the program has been given (if a directory separatory has been supplied). 699 if search(path_sep, binary): 700 # Test that the binary exists. 701 if not access(binary, F_OK): 702 raise RelaxMissingBinaryError(binary) 703 704 # Test that if the binary is executable. 705 if not access(binary, X_OK): 706 raise RelaxNonExecError(binary) 707 708 # The path to the binary has not been given. 709 else: 710 # Get the PATH environmental variable. 711 path = getenv('PATH') 712 713 # Split PATH by the path separator. 714 path_list = split(path, pathsep) 715 716 # Test that the binary exists within the system path (and exit this function instantly once it has been found). 717 for path in path_list: 718 if access(path + sep + binary, F_OK): 719 return 720 721 # The binary is not located in the system path! 722 raise RelaxNoInPathError(binary)

723 724

725 -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, error=None, error_name=None):

726 """Generator function for reading the spin specific data from file. 727 728 Description 729 =========== 730 731 This function writes a columnar formatted file where each line corresponds to a spin system. 732 Spin identification is either through a spin ID string or through columns containing the 733 molecule name, residue name and number, and/or spin name and number. 734 735 736 @param file: The name of the file to write the data to (or alternatively an already opened file object). 737 @type file: str or file object 738 @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). 739 @type dir: str or None 740 @keyword sep: The column separator which, if None, defaults to whitespace. 741 @type sep: str or None 742 @keyword spin_ids: The list of spin ID strings. 743 @type spin_ids: None or list of str 744 @keyword mol_names: The list of molecule names. 745 @type mol_names: None or list of str 746 @keyword res_nums: The list of residue numbers. 747 @type res_nums: None or list of int 748 @keyword res_names: The list of residue names. 749 @type res_names: None or list of str 750 @keyword spin_nums: The list of spin numbers. 751 @type spin_nums: None or list of int 752 @keyword spin_names: The list of spin names. 753 @type spin_names: None or list of str 754 @keyword force: A flag which if True will cause an existing file to be overwritten. 755 @type force: bool 756 @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. 757 @type data: list or list of lists 758 @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. 759 @type data_name: str or list of str 760 @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. 761 @type error: list or list of lists 762 @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. 763 @type error_name: str or list of str 764 """ 765 766 # Data argument tests. 767 if data: 768 # Data is a list of lists. 769 if isinstance(data[0], list): 770 # Data and data_name don't match. 771 if not isinstance(data_name, list): 772 raise RelaxError("The data_name arg '%s' must be a list as the data argument is a list of lists." % data_name) 773 774 # Error doesn't match. 775 if error and (len(data) != len(error) or len(data[0]) != len(error[0])): 776 raise RelaxError("The data arg:\n%s\n\ndoes not have the same dimensions as the error arg:\n%s." % (data, error)) 777 778 # Data is a simple list. 779 else: 780 # Data and data_name don't match. 781 if not isinstance(data_name, str): 782 raise RelaxError("The data_name arg '%s' must be a string as the data argument is a simple list." % data_name) 783 784 # Error doesn't match. 785 if error and len(data) != len(error): 786 raise RelaxError("The data arg:\n%s\n\ndoes not have the same dimensions as the error arg:\n%s." % (data, error)) 787 788 # Error argument tests. 789 if error: 790 # Error is a list of lists. 791 if isinstance(error[0], list): 792 # Error and error_name don't match. 793 if not isinstance(error_name, list): 794 raise RelaxError("The error_name arg '%s' must be a list as the error argument is a list of lists." % error_name) 795 796 # Error is a simple list. 797 else: 798 # Error and error_name don't match. 799 if not isinstance(error_name, str): 800 raise RelaxError("The error_name arg '%s' must be a string as the error argument is a simple list." % error_name) 801 802 # Number of spins check. 803 args = [spin_ids, mol_names, res_nums, res_names, spin_nums, spin_names] 804 arg_names = ['spin_ids', 'mol_names', 'res_nums', 'res_names', 'spin_nums', 'spin_names'] 805 N = None 806 first_arg = None 807 first_arg_name = None 808 for i in range(len(args)): 809 if isinstance(args[i], list): 810 # First list match. 811 if N == None: 812 N = len(args[i]) 813 first_arg = args[i] 814 first_arg_name = arg_names[i] 815 816 # Length check. 817 if len(args[i]) != N: 818 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]))) 819 820 # Nothing?!? 821 if N == None: 822 raise RelaxError("No spin ID data is present.") 823 824 # Data and error length check. 825 if data and len(data) != N: 826 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))) 827 if error and len(error) != N: 828 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))) 829 830 # The spin arguments. 831 args = [spin_ids, mol_names, res_nums, res_names, spin_nums, spin_names] 832 arg_names = ['spin_id', 'mol_name', 'res_num', 'res_name', 'spin_num', 'spin_name'] 833 834 835 # Init. 836 file_data = [[]] 837 838 # Headers - the spin ID info. 839 for i in range(len(args)): 840 if args[i]: 841 file_data[-1].append(arg_names[i]) 842 843 # Headers - the data. 844 if data: 845 # List of lists. 846 if isinstance(data[0], list): 847 # Loop over the list. 848 for i in range(len(data[0])): 849 # The data. 850 file_data[-1].append(data_name[i]) 851 852 # The error. 853 if error: 854 file_data[-1].append(error_name[i]) 855 856 # Simple list. 857 else: 858 # The data. 859 file_data[-1].append(data_name) 860 861 # The error. 862 if error: 863 file_data[-1].append(error_name) 864 865 # Headers - only errors. 866 elif error: 867 # List of lists. 868 if isinstance(error[0], list): 869 for i in range(len(error[0])): 870 file_data[-1].append(error_name[i]) 871 872 # Simple list. 873 else: 874 file_data[-1].append(error_name) 875 876 877 # Spin specific data. 878 for spin_index in range(N): 879 # Append a new data row. 880 file_data.append([]) 881 882 # The spin ID info. 883 for i in range(len(args)): 884 if args[i]: 885 value = args[i][spin_index] 886 if not isinstance(value, str): 887 value = repr(value) 888 file_data[-1].append(value) 889 890 # The data. 891 if data: 892 # List of lists. 893 if isinstance(data[0], list): 894 # Loop over the list. 895 for i in range(len(data[0])): 896 # The data. 897 file_data[-1].append(repr(data[spin_index][i])) 898 899 # The error. 900 if error: 901 file_data[-1].append(repr(error[spin_index][i])) 902 903 # Simple list. 904 else: 905 # The data. 906 file_data[-1].append(repr(data[spin_index])) 907 908 # The error. 909 if error: 910 file_data[-1].append(repr(error[spin_index])) 911 912 # Only errors. 913 elif error: 914 # List of lists. 915 if isinstance(error[0], list): 916 for i in range(len(error[0])): 917 file_data[-1].append(repr(error[spin_index][i])) 918 919 # Simple list. 920 else: 921 file_data[-1].append(repr(error[spin_index])) 922 923 # No data to write, so do nothing! 924 if file_data == [[]]: 925 return 926 927 # Open the file for writing. 928 file = open_write_file(file_name=file, dir=dir, force=force) 929 930 # The number of rows and columns. 931 num_rows = len(file_data) 932 num_cols = len(file_data[0]) 933 934 # Pretty whitespace formatting. 935 if sep == None: 936 # Determine the maximum column widths for nice whitespace formatting. 937 widths = [] 938 for j in range(num_cols): 939 if j == 0: 940 widths.append(len(file_data[0][j]) + 2) 941 else: 942 widths.append(len(file_data[0][j])) 943 for i in range(num_rows): 944 for j in range(num_cols): 945 size = len(file_data[i][j]) 946 if size > widths[j]: 947 widths[j] = size 948 949 # Convert to format strings. 950 formats = [] 951 for j in range(num_cols): 952 formats.append("%%-%ss" % (widths[j] + 4)) 953 954 # The headings. 955 file.write(formats[0] % ("# " + file_data[0][0])) 956 for j in range(1, num_cols): 957 file.write(formats[j] % file_data[0][j]) 958 file.write('\n') 959 960 # The data. 961 for i in range(1, num_rows): 962 # The row. 963 for j in range(num_cols): 964 file.write(formats[j] % file_data[i][j]) 965 file.write('\n') 966 967 # Non-whitespace formatting. 968 else: 969 # The headings. 970 file.write('#') 971 for j in range(num_cols): 972 # The column separator. 973 if j > 0: 974 file.write(sep) 975 976 # The heading. 977 file.write(file_data[0][j]) 978 file.write('\n') 979 980 # The data. 981 for i in range(1, num_rows): 982 # The row. 983 for j in range(num_cols): 984 # The column separator. 985 if j > 0: 986 file.write(sep) 987 988 # The heading. 989 file.write(file_data[i][j]) 990 file.write('\n')

991 992 993

994 -class DummyFileObject:

995 - def __init__(self):

996 """Set up the dummy object to act as a file object.""" 997 998 # Initialise an object for adding the string from all write calls to. 999 self.data = '' 1000 1001 # Set the closed flag. 1002 self.closed = False

1003 1004

1005 - def close(self):

1006 """A method for 'closing' this object.""" 1007 1008 # Set the closed flag. 1009 self.closed = True

1010 1011

1012 - def write(self, str):

1013 """Mimic the file object write() method so that this class can be used as a file object. 1014 1015 @param str: The string to be written. 1016 @type str: str 1017 """ 1018 1019 # Check if the file is closed. 1020 if self.closed: 1021 raise ValueError('I/O operation on closed file') 1022 1023 # Append the string to the data object. 1024 self.data = self.data + str

1025 1026

1027 - def readlines(self):

1028 """Mimic the file object readlines() method. 1029 1030 This method works even if this dummy file object is closed! 1031 1032 1033 @return: The contents of the file object separated by newline characters. 1034 @rtype: list of str 1035 """ 1036 1037 # Split up the string. 1038 lines = split(self.data, '\n') 1039 1040 # Remove the last line if empty. 1041 if lines[-1] == '': 1042 lines.pop() 1043 1044 # Loop over the lines, re-adding the newline character to match the file object readlines() method. 1045 for i in xrange(len(lines)): 1046 lines[i] = lines[i] + '\n' 1047 1048 # Return the file lines. 1049 return lines

1050 1051 1052

1053 -class SplitIO:

1054 - def __init__(self):

1055 """Class for splitting an IO stream to two outputs."""

1056 1057

1058 - def flush(self):

1059 """Flush all streams.""" 1060 1061 # Call the streams' methods. 1062 self.stream1.flush() 1063 self.stream2.flush()

1064 1065

1066 - def isatty(self):

1067 """Check that both streams are TTYs. 1068 1069 @return: True, only if both streams are TTYs. 1070 @rtype: bool 1071 """ 1072 1073 # Check both streams. 1074 return self.stream1.isatty() & self.stream2.isatty()

1075 1076

1077 - def split(self, stream1, stream2):

1078 """Function for setting the streams.""" 1079 1080 # Arguments. 1081 self.stream1 = stream1 1082 self.stream2 = stream2

1083 1084

1085 - def write(self, text):

1086 """Replacement write function.""" 1087 1088 # Write to stream1. 1089 self.stream1.write(text) 1090 1091 # Write to stream2. 1092 self.stream2.write(text)

1093

Source Code for Module relax_io