B
    b>                 @   s   d 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	m
Z
 y.ddlZejZed dksded	krled
W n ek
r   eddY nX ddddZee	 dZdd Zdd Zdd ZdddZG dd dZG dd dZdS )a  I/O function wrappers for the RDF/CDAO file format.

This is an RDF format that conforms to the Comparative Data Analysis Ontology (CDAO).
See: http://evolutionaryontology.org/cdao

This module requires the librdf Python bindings (http://www.librdf.org)

The CDAOIO.Parser, in addition to parsing text files, can also parse directly
from a triple store that implements the Redland storage interface; similarly,
the CDAOIO.Writer can store triples in a triple store instead of serializing
them to a file.
    N)StringIO)MissingPythonDependencyError)CDAO   )cdao_namespacesresolve_uri)12)z3.0.0z3.1.0z3.2.0z=Support for CDAO tree format requires RDFlib v3.2.1 or later.z-Support for CDAO tree format requires RDFlib.zhttp://www.w3.org/2002/07/owl#z+http://www.w3.org/1999/02/22-rdf-syntax-ns#z%http://www.w3.org/2000/01/rdf-schema#)ZowlZrdfZrdfs   c             C   s   t | tdS )zResolve URI for librdf.)Z
namespaces)r   RDF_NAMESPACES)x r   /lib/python3.7/site-packages/Bio/Phylo/CDAOIO.pyqUri;   s    r   c             C   s   |  ddS )zFormat label for librdf._ )replace)r   r   r   r   format_label@   s    r   c             K   s   t | jf |S )zlIterate over the trees in a CDAO file handle.

    :returns: generator of Bio.Phylo.CDAO.Tree objects.

    )Parserparse)handlekwargsr   r   r   r   I   s    r   Fc             K   s   t | j|fd|i|S )zdWrite a trees in CDAO format to the given file handle.

    :returns: number of trees written.

    plain)Writerwrite)treesr   r   r   r   r   r   r   R   s    r   c               @   s\   e Zd ZdZdddZedd Zdd ZdddZdddZ	dd Z
dddZdd ZdS )r   z&Parse a CDAO tree given a file handle.Nc             C   s"   || _ d| _d| _i | _d| _dS )zInitialize CDAO tree parser.NF)r   graph	node_infochildrenrooted)selfr   r   r   r   __init__b   s
    zParser.__init__c             C   s   t |}| |S )z,Instantiate the class from the given string.)r   )clsZtreetextr   r   r   r   from_stringj   s    zParser.from_stringc             K   s   | j f | |  S )z7Parse the text stream this object was initialized with.)parse_handle_to_graphparse_graph)r    r   r   r   r   r   p   s    zParser.parseFturtlec       	      K   s   | j dkrt | _ | j }x t D ]\}}||| q$W || _d|krT|d }ndtj	| j
jdd }|j| j
||d | j||dS )z,Parse self.handle into RDF model self.model.Nbase_urizfile://\/)fileZpublicIDformat)context)r   rdflibZGraphr   itemsZbindr   ospathabspathr   namer   r   r%   )	r    r   Zparse_formatr,   r   r   kvr'   r   r   r   r$   u   s    


zParser.parse_handle_to_graphc             c   sL   |dkr| j }| j||d x*| jD ] }| |}tj|| jdV  q$W dS )z4Iterate over RDF model yielding CDAO.Tree instances.N)r,   )rootr   )r   get_node_info
tree_rootsparse_childrenr   ZTreer   )r    r   r,   Z	root_nodeclader   r   r   r%      s    
zParser.parse_graphc             C   sb   | j | }i }d|kr"|d |d< d|kr>|d dd|d< d|krR|d |d< tjf |}|S )z2Return a CDAO.Clade object for a given named node.branch_lengthlabelr   r   r2   
confidence)r   r   r   ZClade)r    noderesultr   r9   r   r   r   	new_clade   s    
zParser.new_cladec             C   s  i | _ i | _i | _t | _t | _tddtddtddtddtd	d
tddtddi}x|D ]\}}}t|t|t|  }}}|| jkri | j|< | j| }y|||| < W n tk
r   Y nX |tdkr|tdtdfkr| j	| |tdkrd| j	| qdW x| jD ]}i | j |< | j | }	| j| }
d|
kr| j|
d  }d|kr| j|d  }d|krt
|d |	d< d
|
kr| j|
d
  }d|kr|d |	d< d|
kr|
d }|| jkrg | j|< | j| | qW dS )zGCreate a dictionary containing information about all nodes in the tree.zcdao:has_Parentparentzcdao:belongs_to_Edge_as_Childedgezcdao:has_Annotation
annotationzcdao:has_Valuevaluezcdao:represents_TUtuz
rdfs:labelr;   zcdao:has_Support_Valuer<   zrdf:typezcdao:AncestralNodezcdao:TerminalNodezcdao:has_Rootr:   N)r   Zobj_infor   setZnodesr7   r   strKeyErroraddfloatappend)r    r   r,   Zassignmentssr4   othisr=   r   objrA   rB   rD   r@   r   r   r   r6      sZ    












zParser.get_node_infoc                s:     |}| jkr j| ng } fdd|D |_|S )zTraverse the tree to create a nested clade structure.

        Return a CDAO.Clade, and calls itself recursively for each child,
        traversing the entire tree and creating a nested structure of CDAO.Clade
        objects.
        c                s   g | ]}  |qS r   )r8   ).0Z
child_node)r    r   r   
<listcomp>   s    z)Parser.parse_children.<locals>.<listcomp>)r?   r   clades)r    r=   r9   r   r   )r    r   r8      s    
zParser.parse_children)N)Fr&   N)NN)N)__name__
__module____qualname____doc__r!   classmethodr#   r   r$   r%   r?   r6   r8   r   r   r   r   r   _   s   



Hr   c               @   s8   e Zd ZdZeZdd ZdddZdd	 ZdddZ	d
S )r   z8Based on the writer in Bio.Nexus.Trees (str, to_string).c             C   s"   || _ d| _d| _d| _d| _dS )z.Initialize parameters for writing a CDAO tree.r   N)r   node_counteredge_counter
tu_countertree_counter)r    r   r   r   r   r!   	  s
    zWriter.__init__ Fc             K   s   || _ || _|r"|ds"|d7 }| j}|r:|d|  x(| j D ]\}}|d||f  qFW |d| jd   xN|D ]F}	|  jd7  _d| _|	j	}
| j
|
|	d}x|D ]}| || qW q~W d	S )
z-Write this instance's trees to a file handle.r)   z@base <%s>
z@prefix %s: <%s> .
z<%s> a owl:Ontology .
Zcdaor   ztree%s)r5   N)r   record_complete_ancestryendswithr   r   prefixesr.   rZ   tree_urir9   process_cladeadd_stmt_to_handle)r    r   r_   r\   r   r   r   r3   r4   ZtreeZfirst_clade
statementsstmtr   r   r   r     s"    	

zWriter.writec       
      C   s   g }xt |D ]\}}t|tjrt|}d}xB| j D ]4\}}	||	r:||	d| d}|dkrjd}d}q:W |s~d|kr|	| q|	d|  qt|tj
r|	|  q|	t| qW |d	d
|  dS )zAdd URI prefix to handle.Fz%s:r   zrdf:typeaT:z<%s>z%s .
r   N)	enumerate
isinstancer-   URIRefrF   r^   r.   
startswithr   rJ   LiteralZn3r   join)
r    r   rc   Zstmt_stringsnpartZnode_uriZchangedprefixurir   r   r   ra   3  s$    
zWriter.add_stmt_to_handleNc             #   s  |  j d7  _ dt| j t  _|r:|j|jg  _ng  _dd dd d}g }|dk	r| jrrd	nd
}||d|f|d jfg7 }y
|j}W n tk
r   g }Y nX x |D ]\}}	|	|||	f qW  j
r|  jd7  _dt| jt }
||
ddf jd|
f|
dtt j
fg7 }y
 j}W n tk
r   g }Y nX x"|D ]\}}	|
||	fV  qW   rdnd}| jd|f jd|fg7 }|dk	r|  jd7  _dt| jt }||ddf|d|f|d|jf|d jf jd|f jd|jf|jd|fg7 }y
 j}W n tk
r   Y n4X |dk	rtj|dd}| jd|fg7 }| jr:t jdkr:| fd d! jD 7 } jdk	rd"t| jt }tj jtdd}||dd#f|d$|f|d%|fg7 }y
 j}W n tk
r   g }Y nX x"|D ]\}}	|||	fV  qW |E dH  y
 j}W n tk
r(   g }Y nX x$|D ]\}}	 j||	fV  q0W   sx& jD ]}| j| dd&E dH  qbW dS )'z9Recursively generate triples describing a tree of clades.r   znode%sc             S   s
   t | S )N)r-   rh   )rK   r   r   r   nUriW  s    z"Writer.process_clade.<locals>.nUric             S   s   t t| S )N)r-   rh   r   )rK   r   r   r   pUri[  s    z"Writer.process_clade.<locals>.pUrir[   Fzcdao:RootedTreezcdao:UnrootedTreezrdf:typezcdao:has_Rootztu%szcdao:TUzcdao:represents_TUz
rdfs:labelzcdao:TerminalNodezcdao:AncestralNodezcdao:belongs_to_TreeNzedge%szcdao:DirectedEdgezcdao:has_Parent_Nodezcdao:has_Child_Nodezcdao:belongs_to_Edge_as_Childzcdao:has_Parentzcdao:belongs_to_Edge_as_Parentz(http://www.w3.org/2001/XMLSchema#decimal)Zdatatypezcdao:has_Support_Valuer   c                s$   g | ]} j d |fqS )zcdao:has_Ancestor)ro   )rO   Zancestor)r9   rp   rq   r   r   rP     s   z(Writer.process_clade.<locals>.<listcomp>zedge_annotation%szcdao:EdgeLengthzcdao:has_Annotationzcdao:has_Value)r@   r5   )rW   rF   zfillZEROESro   Z	ancestorsr   Z
attributesAttributeErrorrJ   r2   rY   r-   rj   r   tu_attributesZis_terminalrX   r<   r\   lenr:   rh   edge_attributesrQ   r`   )r    r9   r@   r5   Ztree_idrb   Z	tree_typeZtree_attributesZ	predicaterN   Ztu_uriru   Z	node_typeZedge_urir<   Zedge_ann_urir:   rw   Zclade_attributesr?   r   )r9   rp   rq   r   r`   N  s    













zWriter.process_clade)r[   FF)NF)
rR   rS   rT   rU   r   r^   r!   r   ra   r`   r   r   r   r   r     s     
r   )F)rU   r/   ior   ZBior   Z	Bio.Phylor   Z	_cdao_owlr   r   r-   __version__ZrdfverImportErrorr   updaters   r   r   r   r   r   r   r   r   r   r   <module>   s6   


		
 &