1#A* ------------------------------------------------------------------- 2#B* This file contains source code for the PyMOL computer program 3#C* Copyright (c) Schrodinger, LLC. 4#D* ------------------------------------------------------------------- 5#E* It is unlawful to modify or remove this copyright notice. 6#F* ------------------------------------------------------------------- 7#G* Please see the accompanying LICENSE file for further information. 8#H* ------------------------------------------------------------------- 9#I* Additional authors of this source file include: 10#-* 11#-* 12#-* 13#Z* ------------------------------------------------------------------- 14 15from __future__ import print_function, absolute_import 16 17if True: 18 import os 19 import sys 20 if sys.version_info[0] == 2: 21 import thread 22 else: 23 import _thread as thread 24 from . import selector 25 import re 26 import copy 27 28 import pymol 29 cmd = sys.modules["pymol.cmd"] 30 from .cmd import _cmd,lock,unlock,Shortcut,QuietException 31 from chempy import io 32 from chempy.pkl import cPickle 33 from chempy.sdf import SDF,SDFRec 34 from .cmd import _feedback,fb_module,fb_mask, \ 35 DEFAULT_ERROR, DEFAULT_SUCCESS, _raising, is_ok, is_error, \ 36 is_list, is_dict, is_tuple, loadable 37 import traceback 38 39 def copy_image(quiet=1,_self=cmd): # incentive feature / proprietary 40 r = DEFAULT_ERROR 41 if _self.is_gui_thread(): 42 r = _self._copy_image(_self,int(quiet)) 43 else: 44 r = _self.do('cmd._copy_image(quiet=%d)'%int(quiet)) 45 if _self._raising(r,_self): raise QuietException 46 return r 47 48 cache_action_dict = { 49 'enable' : 0, 50 'disable' : 1, 51 'read_only' : 2, 52 'clear' : 3, 53 'optimize' : 4, 54 } 55 56 cache_action_sc = Shortcut(cache_action_dict.keys()) 57 58 def cache(action='optimize', scenes='',state=-1, quiet=1, _self=cmd): 59 ''' 60DESCRIPTION 61 62 "cache" manages storage of precomputed results, such as 63 molecular surfaces. 64 65USAGE 66 67 cache action [, scenes [, state ]] 68 69ARGUMENTS 70 71 action = string: enable, disable, read_only, clear, or optimize 72 73 scenes = string: a space-separated list of scene names (default: '') 74 75 state = integer: state index (default: -1) 76 77EXAMPLES 78 79 cache enable 80 cache optimize 81 cache optimize, F1 F2 F5 82 83NOTES 84 85 "cache optimize" will iterate through the list of scenes provided 86 (or all defined scenes), compute any missing surfaces, and store 87 them in the cache for later reuse. 88 89PYMOL API 90 91 cmd.cache(string action, string scenes, int state, int quiet) 92 93 ''' 94 95 r = DEFAULT_ERROR 96 action = cache_action_dict[cache_action_sc.auto_err(str(action),'action')] 97 quiet = int(quiet) 98 if action == 0: # enable 99 r = _self.set('cache_mode',2,quiet=quiet) 100 elif action == 1: # disable 101 r =_self.set('cache_mode',0,quiet=quiet) 102 elif action == 2: # read_only 103 r =_self.set('cache_mode',1,quiet=quiet) 104 elif action == 3: # clear 105 r =_self._cache_clear(_self=_self) 106 elif action == 4: # optimize 107 r = DEFAULT_SUCCESS 108 _self._cache_mark(_self=_self) 109 cur_scene = _self.get('scene_current_name') 110 cache_max = int(_self.get('cache_max')) 111 if cache_max>0: 112 # allow double memory for an optimized cache 113 _self.set('cache_max',cache_max*2) 114 scenes = str(scenes) 115 scene_list = scenes.split() 116 cache_mode = int(_self.get('cache_mode')) 117 _self.set('cache_mode',2) 118 if not len(scene_list): 119 scene_list = _self.get_scene_list() 120 for scene in scene_list: 121 scene = scene.strip() 122 if not quiet: 123 print(" cache: optimizing scene '%s'."%scene) 124 _self.scene(scene,animate=0) 125 _self.rebuild() 126 _self.refresh() 127 if len(cur_scene): 128 _self.scene(cur_scene,animate=0) 129 else: 130 scene_list = _self.get_scene_list() 131 if len(scene_list): 132 _self.scene(scene_list[0],animate=0) 133 else: 134 if not quiet: 135 print(" cache: no scenes defined -- optimizing current display.") 136 _self.rebuild() 137 _self.refresh() 138 usage = _self._cache_purge(-1,_self=_self) 139 if cache_mode: 140 _self.set('cache_mode',cache_mode) 141 else: 142 _self.set('cache_mode',2) # hmm... could use 1 here instead. 143 _self.set('cache_max',cache_max) # restore previous limits 144 if not quiet: 145 print(" cache: optimization complete (~%0.1f MB)."%(usage*4/1000000.0)) 146 try: 147 _self.lock(_self) 148 finally: 149 _self.unlock(r,_self) 150 if _self._raising(r,_self): raise QuietException 151 return r 152 153 _resn_to_aa = { 154 'ALA' : 'A', 155 'CYS' : 'C', 156 'ASP' : 'D', 157 'GLU' : 'E', 158 'PHE' : 'F', 159 'GLY' : 'G', 160 'HIS' : 'H', 161 'ILE' : 'I', 162 'LYS' : 'K', 163 'LEU' : 'L', 164 'MET' : 'M', 165 'ASN' : 'N', 166 'PRO' : 'P', 167 'GLN' : 'Q', 168 'ARG' : 'R', 169 'SER' : 'S', 170 'THR' : 'T', 171 'VAL' : 'V', 172 'TRP' : 'W', 173 'TYR' : 'Y', 174 # RNA 175 'A' : 'A', 176 'U' : 'U', 177 'G' : 'G', 178 'C' : 'C', 179 # DNA 180 'DA' : 'A', 181 'DT' : 'T', 182 'DG' : 'G', 183 'DC' : 'C', 184 } 185 186 def get_fastastr(selection="all", state=-1, quiet=1, key='', _self=cmd): 187 ''' 188DESCRIPTION 189 190 API only. Get protein and nucleic acid sequences in fasta format. 191 192 Used for saving: 193 PyMOL> save foo.fasta 194 195 New in PyMOL 2.2: 196 - chain specific keys (key argument) 197 - nucleic acid support 198 199ARGUMENTS 200 201 selection = str: atom selection (reduced to "guide & alt +A") {default: all} 202 203 state = int: (only used if state > 0) 204 205 quiet = 0/1: UNUSED 206 207 key = str: python expression {default: model + "_" + chain} 208 Use key=model to get the old (non-chain specific) behavior 209 ''' 210 import textwrap 211 import collections 212 213 seq = collections.OrderedDict() 214 lines = [] 215 216 selection = "(" + selection + ") & guide & alt +A" 217 if int(state) > 0: 218 selection += ' & state {}'.format(state) 219 220 if not key: 221 key = 'model + "_" + chain' 222 # for discrete objects: state specific keys 223 key += ' + (":{}".format(state) if state else "")' 224 225 _self.iterate(selection, 226 "seq.setdefault((" + key + "),[]).append(resn)", 227 space={'seq': seq, 'str': str}) 228 229 for key, resn_list in seq.items(): 230 cur_seq = ''.join(_resn_to_aa.get(x, '?') for x in resn_list) 231 lines.append(">{}".format(key)) 232 lines.extend(textwrap.wrap(cur_seq, 70)) 233 234 if lines: 235 lines.append('') # final newline 236 return '\n'.join(lines) 237 238 def get_pdbstr(selection="all", state=-1, ref='', ref_state=-1, quiet=1, _self=cmd): 239 ''' 240DESCRIPTION 241 242 "get_pdbstr" in an API-only function which returns a pdb 243 corresponding to the atoms in the selection provided and that are 244 present in the indicated state 245 246PYMOL API 247 248 cmd.get_pdbstr(string selection, int state) 249 250NOTES 251 252 "state" is a 1-based state index for the object. 253 254 if state is -1, then current state is used. 255 256 if state is 0, then all states are saved. 257 258 ''' 259 return get_str('pdb', selection, state, ref, ref_state, -1, quiet, _self) 260 261 def _get_dump_str(obj): 262 if is_list(obj): 263 list = map(_get_dump_str,obj) 264 result = "[ " + ",\n".join(list) + " ] " 265 elif is_dict(obj): 266 list = [] 267 for key in obj.keys(): 268 list.append( _get_dump_str(key)+" : "+_get_dump_str(obj[key]) ) 269 result = "{ " + ",\n".join(list) + " } " 270 elif is_tuple(obj): 271 list = map(_get_dump_str,obj) 272 result = "( " + ",\n".join(list) + " ) " 273 else: 274 result = str(obj) 275 return result 276 277 def _session_convert_legacy(session, version, _self=cmd): 278 ''' 279 Convert session contents to be compatible with previous PyMOL versions 280 ''' 281 if version >= _self.get_version()[1]: 282 return 283 284 print(" Applying pse_export_version=%.3f compatibility" % (version)) 285 286 def bitmaskToList(mask): 287 if not isinstance(mask, int): 288 return mask 289 r = [] 290 while (mask >> len(r)): 291 r.append(1 if ((mask >> len(r)) & 1) else 0) 292 return r 293 294 def convert_settings(settings_list): 295 if not settings_list: 296 return 297 298 # get index -> setting dict 299 settings = dict((s[0], s) for s in settings_list if s is not None) 300 301 # changed type 302 cast_fn = { 303 3: float, 304 4: lambda v: list(_self.get_color_tuple(v)), 305 } 306 for (i, old, new) in [ 307 (6, 4, 5), # bg_rgb 308 (254, 3, 1), # scenes_changed 309 ]: 310 if i in settings and settings[i][1] == new: 311 settings[i][1] = old 312 settings[i][2] = cast_fn[old](settings[i][2]) 313 314 # changed defaults 315 changed = [] 316 if version < 1.76: 317 changed += [ 318 (91, 7, -1), # cartoon_sampling 319 (93, 6., -1.), # cartoon_loop_quality 320 (102, 10., -1.), # cartoon_oval_quality 321 (104, 9., -1.), # cartoon_tube_quality 322 (378, 11., -1.), # cartoon_putty_quality 323 ] 324 if version < 1.5: 325 changed += [ 326 (421, -1, 9), # sphere_mode 327 ] 328 for (i, old, new) in changed: 329 if i in settings and settings[i][2] == new: 330 settings[i][2] = old 331 332 if 'settings' in session: 333 convert_settings(session['settings']) 334 335 # objects 336 for name in session['names']: 337 if name is None: 338 continue 339 340 # spec repOn 341 if version < 1.76 and name[3] is None: 342 name[3] = [0] * 20 343 344 # only continue for objects 345 if name[1] != 0: 346 continue 347 348 if name[4] == 8: # gadget 349 cobject = name[5][0][0] 350 else: 351 cobject = name[5][0] 352 353 # object visRep 354 if version < 1.76: 355 cobject[3] = bitmaskToList(cobject[3]) 356 357 # object settings 358 convert_settings(cobject[8]) 359 360 # molecule 361 if name[4] == 1: 362 # atoms visRep 363 if version < 1.76: 364 for atom in name[5][7]: 365 atom[20] = bitmaskToList(atom[20]) 366 # state settings 367 for cset in name[5][4]: 368 if cset: 369 convert_settings(cset[7]) 370 371 # map 372 elif name[4] == 2: 373 if version < 1.5: 374 # crystal -> [crystal, spacegroup] 375 for state in name[5][2]: 376 if state and state[1]: 377 state[1] = state[1][0] 378 379 # cgo 380 elif name[4] == 6: 381 if version < 1.506: 382 for state in name[5][2]: 383 if len(state) == 1: 384 # std and ray CGO 385 state.append(state[0]) 386 387 def get_session(names='', partial=0, quiet=1, compress=-1, cache=-1, _self=cmd): 388 session = {} 389 r = DEFAULT_SUCCESS 390 cache = int(cache) 391 compress = int(compress) 392 partial = int(partial) 393 394 pse_export_version = round(_self.get_setting_float('pse_export_version'), 4) 395 legacyscenes = (0 < pse_export_version < 1.76) and _self.get_scene_list() 396 397 if sys.version_info[0] > 2: 398 legacypickle = (0 < pse_export_version < 1.9) 399 if legacypickle: 400 print(' Warning: pse_export_version with Python 3 is experimental') 401 cPickle.configure_legacy_dump(legacypickle) 402 403 if legacyscenes: 404 _self.pymol._scene_dict = {} 405 406 scene_current_name = _self.get('scene_current_name') 407 tempname = '_scene_db96724c3cef00875c3bebb4f348f711' 408 _self.scene(tempname, 'store') 409 410 for name in legacyscenes: 411 _self.scene(name, 'recall', animate=0) 412 wizard = _self.get_wizard() 413 message = wizard.message if getattr(wizard, 'from_scene', 0) else None 414 pymol.viewing._legacy_scene(name, 'store', message, _self=_self) 415 416 _self.scene(tempname, 'recall', animate=0) 417 _self.scene(tempname, 'clear') 418 _self.set('scene_current_name', scene_current_name) 419 420 if cache: 421 cache_opt = int(_self.get('session_cache_optimize')) 422 if cache != 0: 423 cache_mode = int(_self.get('cache_mode')) 424 if ((cache_mode > 0) and (cache_opt != 0)) or (cache_opt==1): 425 _self.cache('optimize') 426 for a in _self._pymol._session_save_tasks: 427 if a is None: 428 try: 429 _self.lock(_self) 430 r = _cmd.get_session(_self._COb,session,str(names), 431 int(partial),int(quiet)) 432 finally: 433 _self.unlock(r,_self) 434 try: 435 session['session'] = copy.deepcopy(_self._pymol.session) 436 if cache and hasattr(_self._pymol,'_cache'): 437 session['cache'] = _self._pymol._cache 438 except: 439 traceback.print_exc() 440 else: 441 try: 442 if is_error(a(*(session,), **{'_self':_self})): 443 r = DEFAULT_ERROR 444 except: 445 traceback.print_exc() 446 print("Error: An error occurred when trying to generate session.") 447 print("Error: The resulting session file may be incomplete.") 448 449 if legacyscenes: 450 del session['moviescenes'] 451 session['scene_dict'] = _self.pymol._scene_dict 452 session['scene_order'] = legacyscenes 453 pymol.viewing._legacy_scene('*', 'clear', _self=_self) 454 del _self.pymol._scene_dict 455 456 if is_ok(r): 457 if pse_export_version > 0.0: 458 try: 459 _session_convert_legacy(session, pse_export_version, _self) 460 except Exception as e: 461 print(' Warning: failed to backport session:', e) 462 463 if(compress<0): 464 compress = _self.get_setting_boolean('session_compression') 465 if(compress): 466 import zlib 467 session = zlib.compress(cPickle.dumps(session, 1)) 468 return session 469 elif _self._raising(r,_self): 470 raise QuietException 471 return r 472 473 def _unit2px(value, dpi, unit=''): 474 '''API only. Returns pixel units given a string representation in other units''' 475 if cmd.is_string(value): 476 m = re.search(r'[a-z].*', value, re.I) 477 if m: 478 value, unit = value[:m.start()], m.group(0).lower() 479 480 if not unit or unit == 'px': 481 return int(value) 482 483 upi = {'in': 1.0, 'mm': 25.4, 'cm': 2.54} 484 if unit not in upi: 485 raise pymol.CmdException('unknown unit, supported units are: ' + 486 ', '.join(upi)) 487 488 if dpi < 1: 489 raise pymol.CmdException('dpi > 0 required with unit "%s" ' 490 '(hint: set the "image_dots_per_inch" setting)' % unit) 491 492 return float(value) * dpi / upi[unit] + 0.5 493 494 def png(filename, width=0, height=0, dpi=-1.0, ray=0, 495 quiet=1, prior=0, format=0, _self=cmd): 496 ''' 497DESCRIPTION 498 499 "png" saves a PNG format image file of the current display. 500 501USAGE 502 503 png filename [, width [, height [, dpi [, ray]]]] 504 505ARGUMENTS 506 507 filename = string: file path to be written 508 509 width = integer or string: width in pixels (without units), inches (in) 510 or centimeters (cm). If unit suffix is given, dpi argument is required 511 as well. If only one of width or height is given, the aspect ratio of 512 the viewport is preserved. {default: 0 (current)} 513 514 height = integer or string: height (see width) {default: 0 (current)} 515 516 dpi = float: dots-per-inch {default -1.0 (unspecified)} 517 518 ray = 0 or 1: should ray be run first {default: 0 (no)} 519 520EXAMPLES 521 522 png image.png 523 png image.png, dpi=300 524 png image.png, 10cm, dpi=300, ray=1 525 526NOTES 527 528 PNG is the only image format supported by PyMOL. 529 530SEE ALSO 531 532 mpng, save 533 534PYMOL API 535 536 cmd.png(string filename, int width, int height, float dpi, 537 int ray, int quiet) 538 ''' 539 r = DEFAULT_ERROR 540 prior = int(prior) 541 542 if format == 'png': 543 format = 0 544 545 if prior: 546 # fetch the prior image, without doing any work (fast-path / non-GLUT thread-safe) 547 r = _self._png(str(filename),0,0,float(dpi),0,int(quiet),1, 548 int(format),_self) 549 if r != 1: # no prior image available -- revert to default behavior 550 if prior < 0: # default is to fall back to actual rendering 551 prior = 0 552 if not prior: 553 dpi = float(dpi) 554 if dpi < 0: 555 dpi = _self.get_setting_float('image_dots_per_inch') 556 width = _unit2px(width, dpi) 557 height = _unit2px(height, dpi) 558 559 if _self.is_gui_thread(): 560 r = _self._png(str(filename),int(width),int(height),float(dpi), 561 int(ray),int(quiet),0,int(format),_self) 562 else: 563 r = _self._do("cmd._png('''%s''',%d,%d,%1.6f,%d,%d,%d,%d)"% 564 (filename,width,height,dpi, 565 ray,int(quiet),0,int(format)),_self=_self) 566 if _self._raising(r,_self): raise QuietException 567 return r 568 569 def multisave(filename, pattern="all", state=-1, 570 append=0, format='', quiet=1, _self=cmd): 571 ''' 572DESCRIPTION 573 574 "multisave" will save a multi-entry PDB file. 575 576 Every object in the given selection (pattern) will have a HEADER and a 577 CRYST (if symmetry is defined) record, and is terminated with END. 578 Loading such a multi-entry PDB file into PyMOL will load each entry 579 as a separate object. 580 581 This behavior is different to the "save" command, where a multi-object 582 selection is written "flat" to a PDB file, without HEADER or CRYST 583 records. 584 585ARGUMENTS 586 587 filename = string: file path to be written 588 589 pattern = str: atom selection (before 1.8.4: object name pattern) 590 591 state = int: object state (-1=current, 0=all) {default: -1} 592 593 append = 0/1: append to existing file {default: 0} 594 595 format = str: file format {default: guess from extension, or 'pdb'} 596 ''' 597 from pymol.importing import filename_to_format 598 _, _, format_guessed, zipped = filename_to_format(filename) 599 600 if zipped: 601 raise pymol.CmdException(zipped + ' not supported with multisave') 602 603 if not format: 604 format = format_guessed or 'pdb' 605 606 if format == 'pmo': 607 raise pymol.CmdException('pmo format not supported anymore') 608 609 if format not in ('pdb', 'cif'): 610 raise pymol.CmdException(format + ' format not supported with multisave') 611 612 s = get_str(format, pattern, state, '', -1, 1, quiet, _self) 613 614 if s is None: 615 if _self._raising(): raise QuietException 616 return DEFAULT_ERROR 617 618 filename = _self.exp_path(filename) 619 620 with open(filename, 'a' if int(append) else 'w') as handle: 621 handle.write(s) 622 623 return DEFAULT_SUCCESS 624 625 def assign_atom_types( selection, format = "mol2", state=1, quiet=1, _self=cmd): 626 r = DEFAULT_ERROR 627 try: 628 _self.lock(_self) 629 # format : mol2/sybyl = 1, macromodel/mmd = 2, global setting atom_type_format = 0 630 r = _cmd.assign_atom_types(_self._COb, selection, int(1), int(state-1), quiet) 631 finally: 632 _self.unlock(r,_self) 633 return r 634 635 def get_str(format, selection='(all)', state=-1, ref='', 636 ref_state=-1, multi=-1, quiet=1, _self=cmd): 637 ''' 638DESCRIPTION 639 640 Like "get_bytes" but return a unicode string. 641 ''' 642 assert format not in ('mmtf',), 'binary format, use get_bytes' 643 b = get_bytes(format, selection, state, ref, ref_state, multi, quiet, _self) 644 if b is None: 645 return None 646 return b.decode('utf-8') 647 648 def get_bytes(format, selection='(all)', state=-1, ref='', 649 ref_state=-1, multi=-1, quiet=1, _self=cmd): 650 ''' 651DESCRIPTION 652 653 API-only function which exports the selection to a molecular file 654 format and returns it as a binary ("bytes") string. 655 656ARGUMENTS 657 658 format = str: pdb|cif|sdf|mol|mol2|mae|pqr|xyz 659 660 selection = str: atom selection {default: all} 661 662 state = int: object state (-1=current, 0=all) {default: -1} 663 664 ref = str: object name which defines reference frame {default: } 665 666 ref_state = int: state of ref object {default: -1} 667 668 multi = int: for multi-entry file formats, 0 = single entry, 669 1 = by object, 2 = by object-state, -1 = format default {default: -1} 670 ''' 671 with _self.lockcm: 672 return _cmd.get_str(_self._COb, str(format), str(selection), 673 int(state) - 1, str(ref), int(ref_state), 674 int(multi), int(quiet)) 675 676 if sys.version_info[0] == 2: 677 get_str = get_bytes 678 679 def multifilesave(filename, selection='*', state=-1, format='', ref='', 680 ref_state=-1, quiet=1, _self=cmd): 681 ''' 682DESCRIPTION 683 684 For a selection that spans multiple molecular objects and/or states, 685 save each object and/or state to a separate file. Takes a filename 686 argument with placeholders: 687 688 {name} : object name 689 {state} : state number 690 {title} : state title 691 {num} : file number 692 {} : object name (first) or state (second) 693 694EXAMPLES 695 696 multifilesave /tmp/{name}.pdb 697 multifilesave /tmp/{name}-{state}.cif, state=0 698 multifilesave /tmp/{}-{}.cif, state=0 699 multifilesave /tmp/{}-{title}.sdf, state=0 700 ''' 701 for (fname, osele, ostate) in multifilenamegen( 702 filename, selection, int(state), _self): 703 r = _self.save(fname, osele, ostate, format, ref, ref_state, quiet) 704 return r 705 706 707 def multifilenamegen(filename, selection, state, _self=cmd): 708 '''Given a filename pattern, atom selection and state argument, 709 Generate object-state specific filenames and selections. 710 ''' 711 import string 712 713 filename = _self.exp_path(filename) 714 715 fmt_keys = [v[1] 716 for v in string.Formatter().parse(filename) 717 if v[1] is not None] 718 719 nindexed = fmt_keys.count('') 720 multiobject = nindexed > 0 or 'name' in fmt_keys or 'num' in fmt_keys 721 multistate = nindexed > 1 or 'state' in fmt_keys or 'title' in fmt_keys 722 723 if not (multiobject or multistate): 724 raise ValueError('need one or more of {name}, {num}, {state}, {title}') 725 726 odata = [] 727 for oname in _self.get_object_list(selection): 728 osele = '(%s) & ?%s' % (selection, oname) 729 first = last = state 730 731 if multistate: 732 if state < 0: 733 first = last = _self.get_object_state(oname) 734 735 if first == 0: 736 first = 1 737 last = _self.count_states('%' + oname) 738 739 for ostate in range(first, last + 1): 740 odata.append((oname, osele, ostate)) 741 742 # pad {state} and {num} with zeros 743 swidth = len(str(max(v[2] for v in odata))) 744 nwidth = len(str(len(odata))) 745 filename = filename.replace('{state}', '{state:0%d}' % swidth) 746 filename = filename.replace('{num}', '{num:0%d}' % nwidth) 747 748 for num, (oname, osele, ostate) in enumerate(odata, 1): 749 fname = filename.format(oname, ostate, 750 name=oname, state=ostate, num=num, 751 title=_self.get_title(oname, ostate)) 752 yield fname, osele, ostate 753 754 755 def save(filename, selection='(all)', state=-1, format='', ref='', 756 ref_state=-1, quiet=1, partial=0,_self=cmd): 757 ''' 758DESCRIPTION 759 760 "save" writes content to a file. 761 762USAGE 763 764 save filename [, selection [, state [, format ]]] 765 766ARGUMENTS 767 768 filename = string: file path to be written 769 770 selection = string: atoms to save {default: (all)} 771 772 state = integer: state to save {default: -1 (current state)} 773 774PYMOL API 775 776 cmd.save(string file, string selection, int state, string format) 777 778NOTES 779 780 The file format is automatically chosen if the extesion is one of 781 the supported output formats: pdb, pqr, mol, sdf, pkl, pkla, mmd, out, 782 dat, mmod, cif, pov, png, pse, psw, aln, fasta, obj, mtl, wrl, dae, idtf, 783 or mol2. 784 785 If the file format is not recognized, then a PDB file is written 786 by default. 787 788 For molecular files and where applicable and supported: 789 790 * if state = -1 (default), then only the current state is written. 791 792 * if state = 0, then a multi-state output file is written. 793 794SEE ALSO 795 796 load, get_model 797 ''' 798 quiet = int(quiet) 799 800 # preprocess selection 801 selection = selector.process(selection) 802 # 803 r = DEFAULT_ERROR 804 805 # analyze filename 806 from pymol.importing import filename_to_format, _eval_func 807 _, _, format_guessed, zipped = filename_to_format(filename) 808 filename = _self.exp_path(filename) 809 810 # file format 811 if not format: 812 if not format_guessed: 813 raise pymol.CmdException('Unrecognized file format') 814 format = format_guessed 815 816 # PyMOL session 817 if format in ('pse', 'psw',): 818 _self.set("session_file", 819 # always use unix-like path separators 820 filename.replace("\\", "/"), quiet=1) 821 if not quiet: 822 print(" Save: Please wait -- writing session file...") 823 824 func_type4 = { 825 'mmod': io.mmd.toFile, 826 'pkl': io.pkl.toFile, # binary pickle 827 'pkla': lambda model, filename: io.pkl.toFile(model, filename, bin=0), # ascii pickle 828 } 829 830 contents = None 831 832 if format in savefunctions: 833 # generic forwarding to format specific save functions 834 func = savefunctions[format] 835 func = _eval_func(func) 836 kw = { 837 'filename': filename, 838 'selection': selection, 839 'name': selection, # alt (get_ccp4str) 840 'state': state, 841 'format': format, 842 'ref': ref, 843 'ref_state': ref_state, 844 'quiet': quiet, 845 'partial': partial, 846 '_self': _self, 847 } 848 849 import inspect 850 spec = inspect.getargspec(func) 851 852 if spec.varargs: 853 print('FIXME: savefunctions[%s]: *args' % (format)) 854 855 if not spec.keywords: 856 kw = dict((n, kw[n]) for n in spec.args if n in kw) 857 858 contents = func(**kw) 859 860 if 'filename' in spec.args: 861 # assume function wrote directly to file and returned a status 862 return contents 863 864 elif format in func_type4: 865 func_type4[format](_self.get_model(selection, state, ref, ref_state), filename) 866 r = DEFAULT_SUCCESS 867 else: 868 raise pymol.CmdException('File format not supported for export') 869 870 # function returned sequence of strings or bytes 871 if isinstance(contents, (tuple, list)) and contents: 872 contents = contents[0][:0].join(contents) 873 874 if cmd.is_string(contents): 875 if not isinstance(contents, bytes): 876 contents = contents.encode() 877 878 if zipped == 'gz': 879 import gzip 880 fopen = gzip.open 881 else: 882 fopen = open 883 if zipped == 'bz2': 884 import bz2 885 contents = bz2.compress(contents) 886 887 with fopen(filename, 'wb') as handle: 888 handle.write(contents) 889 r = DEFAULT_SUCCESS 890 891 if _self._raising(r,_self): raise QuietException 892 893 if not quiet: 894 if r == DEFAULT_SUCCESS: 895 print(' Save: wrote "' + filename + '".') 896 else: 897 print(' Save-Error: no file written') 898 899 return r 900 901 def get_cifstr(selection="all", state=-1, quiet=1, _self=cmd): 902 ''' 903DESCRIPTION 904 905 API-only function which returns a mmCIF string. 906 907SEE ALSO 908 909 get_pdbstr 910 ''' 911 return get_str('cif', selection, state, '', -1, -1, quiet, _self) 912 913 def get_xyzstr(selection, state=-1, quiet=1, _self=cmd): 914 return get_str('xyz', selection, state, '', -1, -1, quiet, _self) 915 916 def get_sdfstr(selection, state=-1, ref='', ref_state=-1, quiet=1, _self=cmd): 917 return get_str('sdf', selection, state, ref, ref_state, -1, quiet, _self) 918 919 def get_mol2str(selection, state=-1, ref='', ref_state=-1, quiet=1, _self=cmd): 920 return get_str('mol2', selection, state, ref, ref_state, -1, quiet, _self) 921 922 def get_alnstr(selection, state=-1, quiet=1, _self=cmd): 923 with _self.lockcm: 924 return _cmd.get_seq_align_str(_self._COb, str(selection), 925 int(state)-1, 0, int(quiet)) 926 927 def get_pqrstr(selection, state=-1, ref='', ref_state=-1, quiet=1, _self=cmd): 928 return get_str('pqr', selection, state, ref, ref_state, -1, quiet, _self) 929 930 def get_maestr(selection, state=-1, ref='', ref_state=-1, quiet=1, _self=cmd): 931 return get_str('mae', selection, state, ref, ref_state, -1, quiet, _self) 932 933 def get_ccp4str(name, state=1, quiet=1, format='ccp4', _self=cmd): 934 ftype = getattr(loadable, format, -1) 935 with _self.lockcm: 936 return _cmd.get_ccp4str(_self._COb, str(name), 937 int(state) - 1, int(quiet), ftype) 938 939 def get_psestr(selection, partial, quiet, _self): 940 if '(' in selection: # ignore selections 941 selection = '' 942 session = _self.get_session(selection, partial, quiet) 943 return cPickle.dumps(session, 1) 944 945 def _get_mtl_obj(format, _self): 946 # TODO mtl not implemented, always returns empty string 947 if format == 'mtl': 948 raise pymol.CmdException('.MTL export not implemented') 949 i = {'mtl': 0, 'obj': 1}.get(format) 950 return _self.get_mtl_obj()[i] 951 952 savefunctions = { 953 'cif': get_str, # mmCIF 954 'xyz': get_str, 955 'pdb': get_str, 956 'pqr': get_str, 957 'sdf': get_str, 958 'mol2': get_str, 959 'mae': get_str, 960 'mol': get_str, 961 'mmtf': get_bytes, 962 963 'pse': get_psestr, 964 'psw': get_psestr, 965 966 'fasta': get_fastastr, 967 'aln': get_alnstr, 968 'ccp4': get_ccp4str, 969 'mrc': get_ccp4str, 970 'map': get_ccp4str, 971 972 'png': png, 973 974 # no arguments (some have a "version" argument) 975 'dae': 'pymol.querying:get_collada', 976 'gltf': 'pymol.querying:get_gltf', 977 'wrl': 'pymol.querying:get_vrml', 978 'pov': 'pymol.querying:get_povray', 979 'idtf': 'pymol.querying:get_idtf', 980 'mtl': _get_mtl_obj, # TODO not implemented 981 'obj': _get_mtl_obj, 982 'stl': 'pymol.lazyio:get_stlstr', 983 } 984