# -*- coding: utf-8 -*-
"""Query builder."""
import json
import logging
import random
from collections import UserList
from typing import Any, Dict, List, Set, TextIO, Union
from .constants import SEED_TYPE_ANNOTATION, SEED_TYPE_INDUCTION, SEED_TYPE_NEIGHBORS, SEED_TYPE_SAMPLE
from .selection import get_subgraph
from ...dsl import BaseEntity
from ...struct import union
from ...tokens import parse_result_to_dsl
logger = logging.getLogger(__name__)
SEED_METHOD = 'type'
SEED_DATA = 'data'
MaybeNodeList = Union[BaseEntity, List[BaseEntity], List[Dict]]
[docs]class Seeding(UserList):
"""Represents a container of seeding methods to apply to a network."""
[docs] def append_induction(self, nodes: MaybeNodeList) -> 'Seeding':
"""Add a seed induction method.
:param nodes: A node or list of nodes
:returns: self for fluid API
"""
return self._append_seed_handle_nodes(SEED_TYPE_INDUCTION, nodes)
[docs] def append_neighbors(self, nodes: MaybeNodeList) -> 'Seeding':
"""Add a seed by neighbors.
:param nodes: A node or list of nodes
:returns: self for fluid API
"""
return self._append_seed_handle_nodes(SEED_TYPE_NEIGHBORS, nodes)
[docs] def append_annotation(self, annotation: str, values: Set[str]) -> 'Seeding':
"""Add a seed induction method for single annotation's values.
:param annotation: The annotation to filter by
:param values: The values of the annotation to keep
:returns: self for fluid API
"""
return self._append_seed(
SEED_TYPE_ANNOTATION, {
'annotations': {
annotation: values,
},
},
)
[docs] def append_sample(self, **kwargs) -> 'Seeding':
"""Add seed induction methods.
Kwargs can have ``number_edges`` or ``number_seed_nodes``.
:returns: self for fluid API
"""
data = {
'seed': random.randint(0, 1000000),
}
data.update(kwargs)
return self._append_seed(SEED_TYPE_SAMPLE, data)
def _append_seed(self, seed_type: str, data: Any) -> 'Seeding':
"""Add a seeding method and returns self.
:returns: self for fluid API
"""
self.append({
SEED_METHOD: seed_type,
SEED_DATA: data,
})
return self
def _append_seed_handle_nodes(self, seed_type: str, nodes: MaybeNodeList) -> 'Seeding':
"""Add a seeding method and returns self.
:param seed_type: The seed type
:param nodes: A node or list of nodes
:returns: self for fluid API
"""
return self._append_seed(seed_type, _handle_nodes(nodes))
[docs] def run(self, graph):
"""Seed the graph or return none if not possible.
:type graph: pybel.BELGraph
:rtype: Optional[pybel.BELGraph]
"""
if not self:
logger.debug('no seeding, returning graph: %s', graph)
return graph
subgraphs = []
for seed in self:
seed_method, seed_data = seed[SEED_METHOD], seed[SEED_DATA]
logger.debug('seeding with %s: %s', seed_method, seed_data)
subgraph = get_subgraph(graph, seed_method=seed_method, seed_data=seed_data)
if subgraph is None:
logger.debug('seed returned empty graph: %s', seed)
continue
subgraphs.append(subgraph)
if not subgraphs:
logger.debug('no subgraphs returned')
return
return union(subgraphs)
[docs] def to_json(self) -> List[Dict]:
"""Serialize this seeding container to a JSON object."""
return list(self)
[docs] def dump(self, file, sort_keys: bool = True, **kwargs) -> None:
"""Dump this seeding container to a file as JSON."""
json.dump(self.to_json(), file, sort_keys=sort_keys, **kwargs)
[docs] def dumps(self, sort_keys: bool = True, **kwargs) -> str:
"""Dump this query to a string as JSON."""
return json.dumps(self.to_json(), sort_keys=sort_keys, **kwargs)
[docs] @staticmethod
def from_json(data) -> 'Seeding':
"""Build a seeding container from a JSON list."""
return Seeding(data)
[docs] @staticmethod
def load(file: TextIO) -> 'Seeding':
"""Load a seeding container from a JSON file."""
return Seeding.from_json(json.load(file))
[docs] @staticmethod
def loads(s: str) -> 'Seeding':
"""Load a seeding container from a JSON string."""
return Seeding.from_json(json.loads(s))
def _handle_nodes(nodes: MaybeNodeList) -> List[BaseEntity]:
"""Handle node(s) that might be dictionaries."""
if isinstance(nodes, BaseEntity):
return [nodes]
return [
(
parse_result_to_dsl(node)
if not isinstance(node, BaseEntity) else
node
)
for node in nodes
]