Source code for kgx.cli.cli_utils
import importlib
import os
import sys
from multiprocessing import Pool
from typing import List, Tuple, Any, Optional, Dict, Set
import networkx
import yaml
import kgx
from kgx import PandasTransformer, NeoTransformer, Validator, RdfTransformer
from kgx.config import get_logger
from kgx.operations.graph_merge import merge_all_graphs
from kgx.operations.summarize_graph import summarize_graph
_transformers = {
'tar': kgx.PandasTransformer,
'csv': kgx.PandasTransformer,
'tsv': kgx.PandasTransformer,
'tsv:neo4j': kgx.PandasTransformer,
'nt': kgx.NtTransformer,
'ttl': kgx.RdfTransformer,
'json': kgx.JsonTransformer,
'jsonl': kgx.JsonlTransformer,
'obojson': kgx.ObographJsonTransformer,
# 'rq': kgx.SparqlTransformer,
'owl': kgx.RdfOwlTransformer,
'rsa': kgx.RsaTransformer
}
log = get_logger()
[docs]def get_transformer(file_format: str) -> Any:
"""
Get a Transformer corresponding to a given file format.
.. note::
This method returns a reference to kgx.Transformer class
and not an instance of kgx.Transformer class.
You will have to instantiate the class by calling its constructor.
Parameters
----------
file_format: str
File format
Returns
-------
Any
Reference to kgx.Transformer class corresponding to ``file_format``
"""
t = _transformers.get(file_format)
if not t:
raise TypeError(f"format '{file_format}' is not a supported file type.")
return t
[docs]def get_file_types() -> Tuple:
"""
Get all file formats supported by KGX.
Returns
-------
Tuple
A tuple of supported file formats
"""
return tuple(_transformers.keys())
[docs]def graph_summary(inputs: List[str], input_format: str, input_compression: Optional[str], output: Optional[str]) -> Dict:
"""
Loads and summarizes a knowledge graph from a set of input files.
Parameters
----------
inputs: List[str]
Input file
input_format: str
Input file format
input_compression: Optional[str]
The input compression type
output: Optional[str]
Where to write the output (stdout, by default)
Returns
-------
Dict
A dictionary with the graph stats
"""
transformer = get_transformer(input_format)()
for file in inputs:
transformer.parse(file, input_format=input_format, compression=input_compression)
stats = summarize_graph(transformer.graph)
if output:
WH = open(output, 'w')
WH.write(yaml.dump(stats))
else:
print(yaml.dump(stats))
return stats
[docs]def validate(inputs: List[str], input_format: str, input_compression: Optional[str], output: Optional[str]) -> List:
"""
Run KGX validator on an input file to check for Biolink Model compliance.
Parameters
----------
inputs: List[str]
Input files
input_format: str
The input format
input_compression: Optional[str]
The input compression type
output: Optional[str]
Path to output file (stdout, by default)
Returns
-------
List
Returns a list of errors, if any
"""
transformer = get_transformer(input_format)()
for file in inputs:
transformer.parse(file, input_format=input_format, compression=input_compression)
validator = Validator()
errors = validator.validate(transformer.graph)
if output:
validator.write_report(errors, open(output, 'w'))
else:
validator.write_report(errors, sys.stdout)
return errors
[docs]def neo4j_download(uri: str, username: str, password: str, output: str, output_format: str, output_compression: Optional[str], node_filters: Optional[Tuple] = None, edge_filters: Optional[Tuple] = None) -> kgx.Transformer:
"""
Download nodes and edges from Neo4j database.
Parameters
----------
uri: str
Neo4j URI. For example, https://localhost:7474
username: str
Username for authentication
password: str
Password for authentication
output: str
Where to write the output (stdout, by default)
output_format: Optional[str]
The output type (``tsv``, by default)
output_compression: Optional[str]
The output compression type
node_filters: Optional[Tuple]
Node filters
edge_filters: Optional[Tuple]
Edge filters
Returns
-------
kgx.Transformer
The NeoTransformer
"""
transformer = NeoTransformer(uri=uri, username=username, password=password)
if node_filters:
for n in node_filters:
transformer.set_node_filter(n[0], n[1])
if edge_filters:
for e in edge_filters:
transformer.set_edge_filter(e[0], e[1])
transformer.load()
if not output_format:
output_format = 'tsv'
output_transformer = get_transformer(output_format)(transformer.graph)
output_transformer.save(output, output_format=output_format)
return output_transformer
[docs]def neo4j_upload(inputs: List[str], input_format: str, input_compression: Optional[str], uri: str, username: str, password: str, node_filters: Optional[Tuple] = None, edge_filters: Optional[Tuple] = None) -> kgx.Transformer:
"""
Upload a set of nodes/edges to a Neo4j database.
Parameters
----------
inputs: List[str]
A list of files that contains nodes/edges
input_format: str
The input format
input_compression: Optional[str]
The input compression type
uri: str
The full HTTP address for Neo4j database
username: str
Username for authentication
password: str
Password for authentication
node_filters: Optional[Tuple]
Node filters
edge_filters: Optional[Tuple]
Edge filters
Returns
-------
kgx.Transformer
The NeoTransformer
"""
transformer = get_transformer(input_format)()
for file in inputs:
transformer.parse(file, input_format=input_format, compression=input_compression)
if node_filters:
for n in node_filters:
transformer.set_node_filter(n[0], n[1])
if edge_filters:
for e in edge_filters:
transformer.set_edge_filter(e[0], e[1])
neo_transformer = NeoTransformer(transformer.graph, uri=uri, username=username, password=password)
neo_transformer.save()
return neo_transformer
[docs]def transform(inputs: Optional[List[str]], input_format: Optional[str] = None, input_compression: Optional[str] = None, output: Optional[str] = None, output_format: Optional[str] = None, output_compression: Optional[str] = None, node_filters: Optional[Tuple] = None, edge_filters: Optional[Tuple] = None, transform_config: str = None, source: Optional[List] = None, destination: Optional[List] = None, processes: int = 1) -> None:
"""
Transform a Knowledge Graph from one serialization form to another.
Parameters
----------
inputs: Optional[List[str]]
A list of files that contains nodes/edges
input_format: Optional[str]
The input format
input_compression: Optional[str]
The input compression type
output: Optional[str]
The output file
output_format: Optional[str]
The output format
output_compression: Optional[str]
The output compression type
node_filters: Optional[Tuple]
Node filters
edge_filters: Optional[Tuple]
Edge filters
transform_config: Optional[str]
The transform config YAML
source: Optional[List]
A list of source to load from the YAML
destination: Optional[List]
A list of destination to write to, as defined in the YAML
processes: int
Number of processes to use
"""
if transform_config and inputs:
raise ValueError("Can accept either --transform-config OR inputs, not both")
output_directory = 'output'
if transform_config:
node_properties = []
predicate_mappings = {}
curie_map = {}
property_types = {}
checkpoint = False
cfg = yaml.load(open(transform_config), Loader=yaml.FullLoader)
if 'configuration' in cfg:
if 'checkpoint' in cfg['configuration'] and cfg['configuration']['checkpoint'] is not None:
checkpoint = cfg['configuration']['checkpoint']
if 'node_properties' in cfg['configuration'] and cfg['configuration']['node_properties']:
node_properties = cfg['configuration']['node_properties']
if 'predicate_mappings' in cfg['configuration'] and cfg['configuration']['predicate_mappings']:
predicate_mappings = cfg['configuration']['predicate_mappings']
if 'curie_map' in cfg['configuration'] and cfg['configuration']['curie_map']:
curie_map = cfg['configuration']['curie_map']
if 'property_types' in cfg['configuration'] and cfg['configuration']['property_types']:
property_types = cfg['configuration']['property_types']
if 'output_directory' in cfg['configuration'] and cfg['configuration']['output_directory']:
output_directory = cfg['configuration']['output_directory']
if not output_directory.startswith(os.path.sep):
# relative path
output_directory = f"{os.path.abspath(os.path.dirname(transform_config))}{os.path.sep}{output_directory}"
if not source:
source = cfg['transform']['source'].keys()
for s in source:
source_properties = cfg['transform']['source'][s]
if source_properties['input']['format'] in get_file_types():
for f in source_properties['input']['filename']:
if not os.path.exists(f):
raise FileNotFoundError(f"Filename '{f}' for source '{s}' does not exist!")
elif not os.path.isfile(f):
raise FileNotFoundError(f"Filename '{f}' for source '{s}' is not a file!")
source_to_parse = {}
for key, val in cfg['transform']['source'].items():
if key in source:
source_to_parse[key] = val
results = []
pool = Pool(processes=processes)
for k, v in source_to_parse.items():
log.info(f"Spawning process for '{k}'")
result = pool.apply_async(transform_source, (k, v, output_directory, curie_map, node_properties, predicate_mappings, property_types, checkpoint, False))
results.append(result)
pool.close()
pool.join()
else:
source_dict: Dict = {
'input': {
'format': input_format,
'compression': input_compression,
'filename': inputs,
},
'output': {
'format': output_format,
'compression': output_compression,
'filename': output
}
}
transform_source(None, source_dict, None)
[docs]def merge(merge_config: str, source: Optional[List] = None, destination: Optional[List] = None, processes: int = 1) -> networkx.MultiDiGraph:
"""
Load nodes and edges from files and KGs, as defined in a config YAML, and merge them into a single graph.
The merged graph can then be written to a local/remote Neo4j instance OR be serialized into a file.
Parameters
----------
merge_config: str
Merge config YAML
source: Optional[List]
A list of source to load from the YAML
destination: Optional[List]
A list of destination to write to, as defined in the YAML
processes: int
Number of processes to use
Returns
-------
networkx.MultiDiGraph
The merged graph
"""
with open(merge_config, 'r') as YML:
cfg = yaml.load(YML, Loader=yaml.FullLoader)
node_properties = []
predicate_mappings = {}
curie_map = {}
property_types = {}
output_directory = 'output'
checkpoint = False
if 'configuration' in cfg:
if 'checkpoint' in cfg['configuration'] and cfg['configuration']['checkpoint'] is not None:
checkpoint = cfg['configuration']['checkpoint']
if 'node_properties' in cfg['configuration'] and cfg['configuration']['node_properties']:
node_properties = cfg['configuration']['node_properties']
if 'predicate_mappings' in cfg['configuration'] and cfg['configuration']['predicate_mappings']:
predicate_mappings = cfg['configuration']['predicate_mappings']
if 'curie_map' in cfg['configuration'] and cfg['configuration']['curie_map']:
curie_map = cfg['configuration']['curie_map']
if 'property_types' in cfg['configuration'] and cfg['configuration']['property_types']:
property_types = cfg['configuration']['property_types']
if 'output_directory' in cfg['configuration'] and cfg['configuration']['output_directory']:
output_directory = cfg['configuration']['output_directory']
if not output_directory.startswith(os.path.sep):
# relative path
output_directory = f"{os.path.abspath(os.path.dirname(merge_config))}{os.path.sep}{output_directory}"
if not source:
source = cfg['merged_graph']['source'].keys()
if not destination:
destination = cfg['merged_graph']['destination'].keys()
for s in source:
source_properties = cfg['merged_graph']['source'][s]
if source_properties['input']['format'] in get_file_types():
for f in source_properties['input']['filename']:
if not os.path.exists(f):
raise FileNotFoundError(f"Filename '{f}' for source '{s}' does not exist!")
elif not os.path.isfile(f):
raise FileNotFoundError(f"Filename '{f}' for source '{s}' is not a file!")
sources_to_parse = {}
for key in cfg['merged_graph']['source']:
if key in source:
sources_to_parse[key] = cfg['merged_graph']['source'][key]
results = []
pool = Pool(processes=processes)
for k, v in sources_to_parse.items():
log.info(f"Spawning process for '{k}'")
result = pool.apply_async(parse_source, (k, v, output_directory, curie_map, node_properties, predicate_mappings, checkpoint))
results.append(result)
pool.close()
pool.join()
graphs = [r.get() for r in results]
merged_graph = merge_all_graphs(graphs)
if 'name' in cfg['merged_graph']:
merged_graph.name = cfg['merged_graph']['name']
if 'operations' in cfg['merged_graph']:
apply_operations(cfg['merged_graph'], merged_graph)
destination_to_write: Dict[str, Dict] = {}
for d in destination:
if d in cfg['merged_graph']['destination']:
destination_to_write[d] = cfg['merged_graph']['destination'][d]
else:
raise KeyError(f"Cannot find destination '{d}' in YAML")
# write the merged graph
if destination_to_write:
for key, destination_info in destination_to_write.items():
log.info(f"Writing merged graph to {key}")
if destination_info['format'] == 'neo4j':
destination_transformer = NeoTransformer(
source_graph=merged_graph,
uri=destination_info['uri'],
username=destination_info['username'],
password=destination_info['password']
)
destination_transformer.save()
elif destination_info['format'] in get_file_types():
destination_transformer = get_transformer(destination_info['format'])(merged_graph)
filename = destination_info['filename']
if isinstance(filename, list):
filename = filename[0]
destination_filename = f"{output_directory}/{filename}"
if destination_info['format'] == 'nt' and isinstance(destination_transformer, RdfTransformer):
destination_transformer.set_predicate_mapping(predicate_mappings)
destination_transformer.set_property_types(property_types)
compression = destination_info['compression'] if 'compression' in destination_info else None
destination_transformer.save(
filename=destination_filename,
output_format=destination_info['format'],
compression=compression
) # type: ignore
else:
log.error(f"type {destination_info['format']} not yet supported for KGX merge operation.")
else:
log.warning(f"No destination provided in {merge_config}. The merged graph will not be persisted.")
return merged_graph
[docs]def parse_source(key: str, source: dict, output_directory: str, curie_map: Dict[str, str] = None, node_properties: Set[str] = None, predicate_mappings: Dict[str, str] = None, checkpoint: bool = False):
"""
Parse a source from a merge config YAML.
Parameters
----------
key: str
Source key
source: Dict
Source configuration
output_directory: str
Location to write output to
curie_map: Dict[str, str]
Non-canonical CURIE mappings
node_properties: Set[str]
A set of predicates that ought to be treated as node properties (This is applicable for RDF)
predicate_mappings: Dict[str, str]
A mapping of predicate IRIs to property names (This is applicable for RDF)
checkpoint: bool
Whether to serialize each individual source to a TSV
Returns
-------
networkx.MultiDiGraph
Returns a networkx.MultiDiGraph corresponding to the source
"""
log.info(f"Processing source '{key}'")
transformer = parse_source_input(key, source, output_directory, curie_map, node_properties, predicate_mappings, None, checkpoint)
return transformer.graph
[docs]def transform_source(key: Optional[str], source: Dict, output_directory: Optional[str], curie_map: Dict[str, str] = None, node_properties: Set[str] = None, predicate_mappings: Dict[str, str] = None, property_types = None, checkpoint: bool = False, preserve_graph: bool = True) -> networkx.MultiDiGraph:
"""
Transform a source from a transform config YAML.
Parameters
----------
key: Optional[str]
Source key
source: Dict
Source configuration
output_directory: Optional[str]
Location to write output to
curie_map: Dict[str, str]
Non-canonical CURIE mappings
node_properties: Set[str]
A set of predicates that ought to be treated as node properties (This is applicable for RDF)
predicate_mappings: Dict[str, str]
A mapping of predicate IRIs to property names (This is applicable for RDF)
property_types: Dict[str, str]
The xml property type for properties that are other than ``xsd:string``.
Relevant for RDF export.
checkpoint: bool
Whether to serialize each individual source to a TSV
preserve_graph: true
Whether or not to preserve the graph corresponding to the source
Returns
-------
networkx.MultiDiGraph
Returns a networkx.MultiDiGraph corresponding to the source
"""
if not key:
key = os.path.basename(source['input']['filename'][0])
log.info(f"Processing source '{key}'")
output_format = source['output']['format']
output_compression = source['output']['compression'] if 'compression' in source['output'] else None
output_filename = source['output']['filename'] if 'filename' in source['output'] else key
if isinstance(output_filename, list):
output = output_filename[0]
else:
output = output_filename
transformer = parse_source_input(key, source, output_directory, curie_map, node_properties, predicate_mappings, property_types, checkpoint)
if output_directory and not output.startswith(output_directory):
output = os.path.join(output_directory, output)
if output_format == 'neo4j':
output_transformer = NeoTransformer(
source_graph=transformer.graph,
uri=source['output']['uri'],
username=source['output']['username'],
password=source['output']['password']
)
output_transformer.save()
elif output_format in get_file_types():
output_transformer = get_transformer(output_format)(transformer.graph)
if output_format == 'nt' and isinstance(output_transformer, RdfTransformer):
output_transformer.set_property_types(property_types)
output_transformer.save(output, output_format=output_format, compression=output_compression) # type: ignore
else:
raise ValueError(f"type {output_format} not yet supported for output")
if not preserve_graph:
output_transformer.graph.clear()
return output_transformer.graph
[docs]def parse_source_input(key: Optional[str], source: Dict, output_directory: Optional[str], curie_map: Dict[str, str] = None, node_properties: Set[str] = None, predicate_mappings: Dict[str, str] = None, property_types = None, checkpoint: bool = False) -> kgx.Transformer:
"""
Parse a source's input from a transform config YAML.
Parameters
----------
key: Optional[str]
Source key
source: Dict
Source configuration
output_directory: Optional[str]
Location to write output to
curie_map: Dict[str, str]
Non-canonical CURIE mappings
node_properties: Set[str]
A set of predicates that ought to be treated as node properties (This is applicable for RDF)
predicate_mappings: Dict[str, str]
A mapping of predicate IRIs to property names (This is applicable for RDF)
property_types: Dict[str, str]
The xml property type for properties that are other than ``xsd:string``.
Relevant for RDF export.
checkpoint: bool
Whether to serialize each individual source to a TSV
Returns
-------
kgx.Transformer
An instance of kgx.Transformer corresponding to the source format
"""
if not key:
key = os.path.basename(source['input']['filename'][0])
source_name = source['input']['name'] if 'name' in source['input'] else key
input_format = source['input']['format']
input_compression = source['input']['compression'] if 'compression' in source['input'] else None
inputs = source['input']['filename']
filters = source['input']['filters'] if 'filters' in source['input'] and source['input']['filters'] is not None else {}
node_filters = filters['node_filters'] if 'node_filters' in filters else {}
edge_filters = filters['edge_filters'] if 'edge_filters' in filters else {}
operations = source['input']['operations'] if 'operations' in source['input'] and source['input']['filters'] is not None else {}
source_curie_map = source['curie_map'] if 'curie_map' in source and source['curie_map'] is not None else {}
if curie_map:
source_curie_map.update(curie_map)
source_predicate_mappings = source['predicate_mappings'] if 'predicate_mappings' in source and source['predicate_mappings'] is not None else {}
if predicate_mappings:
source_predicate_mappings.update(predicate_mappings)
source_node_properties = source['node_properties'] if 'node_properties' in source and source['node_properties'] is not None else []
if node_properties:
source_node_properties.extend(node_properties)
if input_format in {'nt', 'ttl'}:
# Parse RDF file types
transformer = get_transformer(input_format)(curie_map=source_curie_map)
if predicate_mappings:
transformer.set_predicate_mapping(predicate_mappings)
transformer.graph.name = key
if filters:
apply_filters(transformer, node_filters, edge_filters)
for f in inputs:
transformer.parse(
filename=f,
input_format=input_format,
compression=input_compression,
node_property_predicates=source_node_properties,
provided_by=source_name
)
if operations:
apply_operations(source['input'], transformer.graph)
elif input_format in get_file_types():
# Parse other supported file types
transformer = get_transformer(input_format)()
transformer.graph.name = key
if filters:
apply_filters(transformer, node_filters, edge_filters)
for f in inputs:
transformer.parse(
filename=f,
input_format=input_format,
compression=input_compression,
provided_by=source_name
)
if operations:
apply_operations(source['input'], transformer.graph)
elif input_format == 'neo4j':
# Parse Neo4j
transformer = NeoTransformer(
source_graph=None,
uri=source['uri'],
username=source['username'],
password=source['password']
)
transformer.graph.name = key
if filters:
apply_filters(transformer, node_filters, edge_filters)
transformer.load(provided_by=source_name)
if operations:
apply_operations(source['input'], transformer.graph)
transformer.graph.name = key
else:
raise TypeError(f"type {input_format} not yet supported")
if checkpoint:
log.info(f"Writing checkpoint for source '{key}'")
pt = PandasTransformer(transformer.graph)
checkpoint_output = f"{output_directory}/{key}" if output_directory else key
pt.save(filename=checkpoint_output, output_format='tsv', compression=None)
return transformer
[docs]def apply_filters(transformer: kgx.Transformer, node_filters: Optional[Dict], edge_filters: Optional[Dict]) -> kgx.Transformer:
"""
Apply filters to the given transformer.
Parameters
----------
transformer: kgx.Transformer
The transformer corresponding to the source
node_filters: Optional[Dict]
Node filters
edge_filters: Optional[Dict]
Edge filters
Returns
-------
transformer: kgx.Transformer
The transformer with filters applied
"""
if node_filters:
for k, v in node_filters.items():
transformer.set_node_filter(k, set(v))
if edge_filters:
for k, v in edge_filters.items():
transformer.set_edge_filter(k, set(v))
log.info(f"with node filters: {node_filters}")
log.info(f"with edge filters: {edge_filters}")
return transformer
[docs]def apply_operations(source: dict, graph: networkx.MultiDiGraph) -> networkx.MultiDiGraph:
"""
Apply operations as defined in the YAML.
Parameters
----------
source: dict
The source from the YAML
graph: networkx.MultiDiGraph
The graph corresponding to the source
Returns
-------
networkx.MultiDiGraph
The graph corresponding to the source
"""
operations = source['operations']
for operation in operations:
op_name = operation['name']
op_args = operation['args']
module_name = '.'.join(op_name.split('.')[0:-1])
function_name = op_name.split('.')[-1]
f = getattr(importlib.import_module(module_name), function_name)
log.info(f"Applying operation {op_name} with args: {op_args}")
f(graph, **op_args)
return graph