o
    ॵi                     @   s^   d Z ddlZddlZddlZddlmZ ddlmZ ddlm	Z	 ddl
mZ G dd	 d	ZdS )
z$Library to run HHsearch from Python.    N)Sequence)logging)parsers   )utilsc                   @   s   e Zd ZdZdddedee defddZed	efd
dZ	ed	efddZ
ded	efddZdeded	eej fddZdS )HHSearchz&Python wrapper of the HHsearch binary.i@B )maxseqbinary_path	databasesr   c                C   sJ   || _ || _|| _| jD ]}t|d s"td| td| qdS )aB  Initializes the Python HHsearch wrapper.

        Args:
            binary_path: The path to the HHsearch executable.
            databases: A sequence of HHsearch database paths. This should be the
                common prefix for the database files (i.e. up to but not including
                _hhm.ffindex etc.)
            maxseq: The maximum number of rows in an input alignment. Note that this
                parameter is only supported in HHBlits version 3.1 and higher.

        Raises:
            RuntimeError: If HHsearch binary not found within the path.
        z_*z#Could not find HHsearch database %sz!Could not find HHsearch database N)r	   r
   r   globr   error
ValueError)selfr	   r
   r   database_path r   h/home/ubuntu/.local/lib/python3.10/site-packages/modelscope/models/science/unifold/msa/tools/hhsearch.py__init__   s   
zHHSearch.__init__returnc                 C      dS )Nhhrr   r   r   r   r   output_format;      zHHSearch.output_formatc                 C   r   )Na3mr   r   r   r   r   input_format?   r   zHHSearch.input_formatr   c              	   C   s  t  }tj|d}tj|d}t|d}|| W d   n1 s(w   Y  g }| jD ]}|d || q2| j	d|d|dt
| jg| }td	d
| tj|tjtjd}	t d |	 \}
}|	 }W d   n1 s{w   Y  |rtd|
d|dd df t|}| }W d   n1 sw   Y  W d   |S W d   |S 1 sw   Y  |S )z6Queries the database using HHsearch using a given a3m.z	query.a3mz
output.hhrwNz-dz-iz-oz-maxseqzLaunching subprocess "%s" )stdoutstderrzHHsearch queryz(HHSearch failed:
stdout:
%s

stderr:
%s
zutf-8i )r   tmpdir_managerospathjoinopenwriter
   appendr	   strr   r   info
subprocessPopenPIPEtimingcommunicatewaitRuntimeErrordecoderead)r   r   query_tmp_dir
input_pathhhr_pathfdb_cmddb_pathcmdprocessr   r   retcoder   r   r   r   queryC   sX   








#
##zHHSearch.queryoutput_stringinput_sequencec                 C   s   ~t |S )zAGets parsed template hits from the raw string output by the tool.)r   	parse_hhr)r   r;   r<   r   r   r   get_template_hitsj   s   
zHHSearch.get_template_hitsN)__name__
__module____qualname____doc__r&   r   intr   propertyr   r   r:   r   TemplateHitr>   r   r   r   r   r      s,    
'r   )rB   r   r    r(   typingr   abslr   %modelscope.models.science.unifold.msar    r   r   r   r   r   r   <module>   s   