remove_alt keep=first

psico/editing.py

@@ -254,7 +254,7 @@ def stub2ala(selection='all', quiet=1, *, _self=cmd):
     _self.sort()
 
 
-def remove_alt(selection='all', keep='A', quiet=1, *, _self=cmd):
+def remove_alt(selection='all', keep='first', quiet=1, *, _self=cmd):
     '''
 DESCRIPTION
 
@@ -268,13 +268,39 @@ def remove_alt(selection='all', keep='A', quiet=1, *, _self=cmd):
 
     selection = string: atom selection
 
-    keep = string: AltLoc to keep {default: A}
+    keep = string: AltLoc to keep, or 'first' to keep the first observed AltLoc {default: first}
     '''
+    if keep == "first":
+        return remove_alt_keep_first(selection, quiet=quiet, _self=_self)
+
+    if len(keep) != 1:
+        raise CmdException(
+            f"keep must be 'first' or a single letter, got {keep!r}")
+
     _self.remove('(%s) and not alt +%s' % (selection, keep), quiet=int(quiet))
     _self.alter(selection, '(alt,q)=("",1.0)')
     _self.sort()
 
 
+def remove_alt_keep_first(selection='*', *, quiet=1, _self=cmd):
+    '''
+    Remove alternative location atoms, keep the first observed.
+    '''
+    alts = {}
+    expr = "(alt, q) = callback((model, segi, chain, resi, resn, name), alt)"
+
+    def callback(namekey, alt):
+        return ("#", 0.) if alts.setdefault(namekey, alt) != alt else ("", 1.)
+
+    tmpsele = _self.get_unused_name("_sele")
+    _self.select(tmpsele, selection)
+    try:
+        _self.alter(tmpsele, expr, space={"callback": callback})
+        _self.remove(f"{tmpsele} & not alt ''", quiet=quiet)
+    finally:
+        _self.delete(tmpsele)
+
+
 def _common_ss_alter(selection, ss_dict, ss_map, raw='', *, _self=cmd):
     '''
 DESCRIPTION

tests/test_editing.py

@@ -2,7 +2,7 @@
 import psico.querying
 import os
 import pytest
-from pymol import cmd
+from pymol import cmd, CmdException
 from pytest import approx
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), 'data')
@@ -147,10 +147,20 @@ def test_remove_alt():
     cmd.alter('resn ALA', 'alt="A"')
     cmd.alter('resn CYS', 'alt="B"')
     cmd.create('m2', 'm1')
-    psico.editing.remove_alt('m1')
+    cmd.create('m3', 'm1')
+    cmd.alter('m3', 'alt = {"A": "X", "B": "Y"}[alt]')
+    cmd.create('m4', 'm1 m3')
+    assert cmd.count_atoms('m4') == 42
+    psico.editing.remove_alt('m1', keep='A')
     psico.editing.remove_alt('m2', keep='B')
     assert cmd.count_atoms('m1') == 10
     assert cmd.count_atoms('m2') == 11
+    psico.editing.remove_alt('m3', keep='first')
+    assert cmd.count_atoms('m3') == 21
+    psico.editing.remove_alt('m4', keep='first')
+    assert cmd.count_atoms('m4') == 21
+    with pytest.raises(CmdException, match="single letter"):
+        psico.editing.remove_alt('m1', keep='last')
 
 
 @pytest.mark.exe