U
    ZweL#                     @   s  d dl mZ d dlmZ d dlmZmZmZmZm	Z	 d dl
m
Z
 ddlmZ ddlmZ ddlmZmZ dd	lmZmZ dd
lmZmZ ddlmZ ddde
deeeef ee eee eeef dddZG dd deZG dd deZG dd deZG dd deZ dS )    )	ExitStack)PathLike)UnionBinaryIOOptionalIteratorTuple)xopen   )SequenceRecord)FileFormatError)PairedEndReaderPairedEndWriter)FastaReaderFastqReader)FastaWriterFastqWriter)_open_singleNr)
fileformatmode	qualitiesopener)filesr   r   r   returnc                 G   s   t |dkrb|dkr,|d |d kr,tdd|krFt|| ||dS |dk}t|| |||d	S t |dkrd|krt|d | ||dS |dk}t|d | |||d	S td
dS )z
    Open paired-end reads
       Zwar   r
   z)The paired-end output files are identicalr   )r   r   r   ar   r   r   appendz2_open_paired must be called with one or two files.N)len
ValueErrorTwoFilePairedEndReaderTwoFilePairedEndWriterInterleavedPairedEndReaderInterleavedPairedEndWriter)r   r   r   r   r   r    r%   .lib/python3.8/site-packages/dnaio/pairedend.py_open_paired   sF    
      r'   c                   @   s   e Zd ZdZdZddedeeee	f eeee	f e
e dddZed	d
dZeeeef  d	ddZdd	ddZdd Zdd ZdS )r!   z
    Read paired-end reads from two files (not interleaved)

    While this class can be instantiated directly, the recommended way is to
    use `dnaio.open` with appropriate arguments.
    Tr   Nr   r   r   )file1file2r   c             	   C   sb   || _ t B}|t||||d| _|t||||d| _| j| _W 5 Q R X | jj	| _	d S )N)r   r   r   )
r   r   enter_contextr   reader1reader2pop_allclose_closedelivers_qualities)selfr)   r*   r   r   r   stackr%   r%   r&   __init__B   s    	zTwoFilePairedEndReader.__init__r   c                 C   s   | j j d| j d| j dS )Nz(file1=z, file2=))	__class____name__r,   r-   r2   r%   r%   r&   __repr__V   s    zTwoFilePairedEndReader.__repr__c                 c   s   t | j| jD ]<\}}||s@td|j d|j dddd||fV  qztt| j W n tk
rr   Y nX ztt| j W n tk
r   Y nX | jj	| jj	k rtdddd| jj	| jj	krtdddddS )z
        Iterate over the paired reads.
        Each yielded item is a pair of `SequenceRecord` objects.

        Raises a `FileFormatError` if reads are improperly paired.
        z(Reads are improperly paired. Read name 'z' in file 1 does not match 'z' in file 2.NlinezKReads are improperly paired. There are more reads in file 2 than in file 1.zKReads are improperly paired. There are more reads in file 1 than in file 2.)
zipr,   r-   is_mater   namenextiterStopIterationZnumber_of_records)r2   r1r2r%   r%   r&   __iter__Y   s>    
zTwoFilePairedEndReader.__iter__c                 C   s   |    d S Nr0   r9   r%   r%   r&   r/      s    zTwoFilePairedEndReader.closec                 C   s   | S rF   r%   r9   r%   r%   r&   	__enter__   s    z TwoFilePairedEndReader.__enter__c                 G   s   |    d S rF   r/   )r2   excr%   r%   r&   __exit__   s    zTwoFilePairedEndReader.__exit__r8   
__module____qualname____doc__Zpairedr	   r   strr   r   r   r4   r:   r   r   r   rE   r/   rH   rK   r%   r%   r%   r&   r!   8   s   &r!   c                   @   s   e Zd ZdZdZddedeeee	f e
e dddZed	d
dZeeeef  d	ddZdd	ddZdd Zdd ZdS )r#   z
    Read paired-end reads from an interleaved FASTQ file

    While this class can be instantiated directly, the recommended way is to
    use `dnaio.open` with appropriate arguments.
    Tr   Nr(   )filer   c                C   s<   || _ t||||d}t|ttfs(t|| _| jj| _d S )N)r   r   r   )r   r   
isinstancer   r   AssertionErrorreaderr1   )r2   rQ   r   r   r   rT   r%   r%   r&   r4      s
    z#InterleavedPairedEndReader.__init__r5   c                 C   s   | j j d| j dS N(r6   )r7   r8   rT   r9   r%   r%   r&   r:      s    z#InterleavedPairedEndReader.__repr__c              	   c   s   t | j}|D ]p}zt|}W n, tk
rJ   td|j dd dd Y nX ||sttd|j d|j dd d||fV  qd S )Nz0Interleaved input file incomplete: Last record 'z' has no partner.r;   z#Reads are improperly paired. Name 'z' (first) does not match 'z' (second).)rA   rT   r@   rB   r   r?   r>   )r2   itrC   rD   r%   r%   r&   rE      s"    

z#InterleavedPairedEndReader.__iter__c                 C   s   | j   d S rF   )rT   r/   r9   r%   r%   r&   r/      s    z InterleavedPairedEndReader.closec                 C   s   | S rF   r%   r9   r%   r%   r&   rH      s    z$InterleavedPairedEndReader.__enter__c                 G   s   |    d S rF   rI   r2   argsr%   r%   r&   rK      s    z#InterleavedPairedEndReader.__exit__rL   r%   r%   r%   r&   r#      s   r#   c                   @   s   e Zd ZdZddeddeeeef eeeef e	e e	e
 e
dddZed	d
dZdd	ddZdd	ddZdd Zdd ZdS )r"   z
    Write paired-end reads to two files (not interleaved)

    While this class can be instantiated directly, the recommended way is to
    use `dnaio.open` with appropriate arguments.
    fastqNFr   )r)   r*   r   r   r   c          	   
   C   sj   |rdnd}t  N}|  |  |t|||||d| _|t|||||d| _| j| _W 5 Q R X d S Nr   w)r   r   r   r   )r   r+   r   _writer1_writer2r.   r/   r0   )	r2   r)   r*   r   r   r   r   r   r3   r%   r%   r&   r4      s.    
		zTwoFilePairedEndWriter.__init__r5   c                 C   s   | j j d| j d| j dS )NrV   z, r6   )r7   r8   r]   r^   r9   r%   r%   r&   r:      s    zTwoFilePairedEndWriter.__repr__c                 C   s   | j | | j| d S rF   )r]   writer^   r2   read1read2r%   r%   r&   r_      s    zTwoFilePairedEndWriter.writec                 C   s   |    d S rF   rG   r9   r%   r%   r&   r/      s    zTwoFilePairedEndWriter.closec                 C   s   | S rF   r%   r9   r%   r%   r&   rH      s    z TwoFilePairedEndWriter.__enter__c                 G   s   |    d S rF   rI   rX   r%   r%   r&   rK      s    zTwoFilePairedEndWriter.__exit__)r8   rM   rN   rO   r	   r   rP   r   r   r   boolr4   r:   r_   r/   rH   rK   r%   r%   r%   r&   r"      s"   "r"   c                   @   s   e Zd ZdZddeddeeeef e	e e	e
 e
dddZed	d
dZeeddddZdd	ddZdd Zdd ZdS )r$   z
    Write paired-end reads to an interleaved FASTA or FASTQ file

    While this class can be instantiated directly, the recommended way is to
    use `dnaio.open` with appropriate arguments.
    rZ   NFr   )rQ   r   r   r   c                C   s:   |rdnd}t |||||d}t|ttfs0t|| _d S r[   )r   rR   r   r   rS   _writer)r2   rQ   r   r   r   r   r   writerr%   r%   r&   r4     s    	    z#InterleavedPairedEndWriter.__init__r5   c                 C   s   | j j d| j dS rU   )r7   r8   rd   r9   r%   r%   r&   r:     s    z#InterleavedPairedEndWriter.__repr__)ra   rb   r   c                 C   s   | j | | j | d S rF   )rd   r_   r`   r%   r%   r&   r_     s    z InterleavedPairedEndWriter.writec                 C   s   | j   d S rF   )rd   r/   r9   r%   r%   r&   r/     s    z InterleavedPairedEndWriter.closec                 C   s   | S rF   r%   r9   r%   r%   r&   rH     s    z$InterleavedPairedEndWriter.__enter__c                 G   s   |    d S rF   rI   rX   r%   r%   r&   rK   #  s    z#InterleavedPairedEndWriter.__exit__)r8   rM   rN   rO   r	   r   rP   r   r   r   rc   r4   r:   r   r_   r/   rH   rK   r%   r%   r%   r&   r$      s    r$   )!
contextlibr   osr   typingr   r   r   r   r   r	   Z_corer   
exceptionsr   Z
interfacesr   r   Zreadersr   r   Zwritersr   r   Z	singleendr   rP   rc   r'   r!   r#   r"   r$   r%   r%   r%   r&   <module>   s0   
)Q8<