def restrain_ss(session, residues, ss_type):
    from chimerax.atomic import Residue
    structures = residues.unique_structures
    for s in structures:
        sresidues = s.residues.intersect(residues)
        strands = sresidues[sresidues.ss_types==Residue.SS_STRAND]
        helices = sresidues[sresidues.ss_types==Residue.SS_HELIX]
        from chimerax.isolde.restraints.restraint_utils import restrain_secondary_structure
        if ss_type == 'current':
            restrain_secondary_structure(session, strands, 'strand')
            restrain_secondary_structure(session, helices, 'helix')
        if ss_type == 'strand':
            restrain_secondary_structure(session, sresidues, 'strand')
        if ss_type == 'helix':
            restrain_secondary_structure(session, sresidues, 'helix')

from chimerax.core.commands import CmdDesc, register, EnumOf
from chimerax.atomic import ResiduesArg


desc = CmdDesc(
    synopsis='Add backbone distance and torsion restraints to restrain residues to idealised secondary structure geometry',
    required=[
        ('residues', ResiduesArg),
        ('ss_type', EnumOf(('strand', 'helix', 'current')))
    ],
)
register('isolde restrain ss', desc, restrain_ss, logger=session.logger)