Skip to content

Commit

Permalink
Add py2neo version 3 URI compatibility (#8)
Browse files Browse the repository at this point in the history
Closes #4
Closes #8
Closes #6
  • Loading branch information
veleritas authored and dhimmel committed Mar 2, 2017
1 parent 48b65b0 commit 7fe01f1
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 6 deletions.
40 changes: 34 additions & 6 deletions hetio/neo4j.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,18 +4,29 @@
import random
import re

from operator import or_
from functools import reduce

import py2neo
import py2neo.packages.httpstream
import pandas

from tqdm import tqdm

PY2NEO_VER = int(py2neo.__version__[0])

import hetio.hetnet

# Avoid SocketError
py2neo.packages.httpstream.http.socket_timeout = 1e8

def export_neo4j(graph, uri, node_queue=100, edge_queue=100):
def export_neo4j(graph, uri, node_queue=100, edge_queue=100, show_progress=False):
"""Export hetnet to neo4j"""
db_graph = py2neo.Graph(uri)

if isinstance(uri, py2neo.Graph):
db_graph = uri
else:
db_graph = py2neo.Graph(uri)

# Delete all existing nodes
db_graph.delete_all()
Expand All @@ -28,18 +39,30 @@ def export_neo4j(graph, uri, node_queue=100, edge_queue=100):
if 'name' not in db_graph.schema.get_indexes(label):
db_graph.schema.create_index(label, 'name')


# Create nodes
creator = Creator(db_graph, node_queue)
for node in graph.get_nodes():

queue = graph.get_nodes()
if show_progress:
queue = tqdm(queue, total=graph.n_nodes, desc="Importing nodes")

for node in queue:
label = as_label(node.metanode)
data = sanitize_data(node.data)
neo_node = py2neo.Node(label, identifier=node.identifier, name=node.name, **data)
creator.append(neo_node)
creator.create()


# Create edges
creator = Creator(db_graph, edge_queue)
for edge in graph.get_edges(exclude_inverts=True):

queue = graph.get_edges(exclude_inverts=True)
if show_progress:
queue = tqdm(queue, total=graph.n_edges, desc = "Importing edges")

for edge in queue:
metaedge = edge.metaedge
rel_type = as_type(metaedge)
source_label = as_label(metaedge.source)
Expand All @@ -53,7 +76,6 @@ def export_neo4j(graph, uri, node_queue=100, edge_queue=100):

return db_graph


class Creator(list):
"""Batch creation of py2neo objects"""

Expand All @@ -71,7 +93,13 @@ def append(self, x):
def create(self):
if not self:
return
self.db_graph.create(*self)

# http://stackoverflow.com/questions/37676731/typeerror-create-takes-2-positional-arguments-but-4-were-given
if PY2NEO_VER >= 3:
self.db_graph.create(reduce(or_, self))
else:
self.db_graph.create(*self)

self.n_created += len(self)
#print('{} nodes created\r'.format(self.n_created), end='')
del self[:]
Expand Down
2 changes: 2 additions & 0 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,4 +23,6 @@
'Programming Language :: Python :: 3.4',
'Programming Language :: Python :: 3.5',
],

install_requires=['tqdm']
)

0 comments on commit 7fe01f1

Please sign in to comment.