Skip to content

Commit c4e18b1

Browse files
Can now export a NetworkX graph with attributes and version bump
1 parent 03314cc commit c4e18b1

File tree

3 files changed

+194
-42
lines changed

3 files changed

+194
-42
lines changed

django_postgresql_dag/exceptions.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,3 +8,19 @@ class NodeNotReachableException(Exception):
88

99
pass
1010

11+
12+
class GraphModelsCannotBeParsedException(Exception):
13+
"""
14+
Exception for node distance and path
15+
"""
16+
17+
pass
18+
19+
20+
class IncorrectUsageException(Exception):
21+
"""
22+
Exception for node distance and path
23+
"""
24+
25+
pass
26+

django_postgresql_dag/transformations.py

Lines changed: 177 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,16 @@
33
django-postgresql-dag to alternate formats.
44
"""
55

6+
from django.core.exceptions import FieldDoesNotExist
67
from django.db.models import Case, When
8+
from django.db.models.fields import DateTimeField, UUIDField
9+
from django.db.models.fields.files import ImageField, FileField
10+
from django.db.models.fields.related import ManyToManyField
711

12+
from .exceptions import GraphModelsCannotBeParsedException, IncorrectUsageException
13+
14+
from itertools import chain
815
import networkx as nx
9-
import pandas as pd
1016

1117

1218
def _filter_order(queryset, field_names, values):
@@ -30,54 +36,184 @@ def _filter_order(queryset, field_names, values):
3036
return queryset.filter(**filter_condition).order_by(order_by)
3137

3238

33-
def rawqueryset_to_values_list(rawqueryset):
34-
"""Returns a list of lists of each instance"""
35-
columns = rawqueryset.columns
36-
for row in rawqueryset:
37-
yield tuple(getattr(row, column) for column in columns)
39+
def get_queryset_characteristics(queryset):
40+
"""
41+
Returns a tuple of the node & edge model classes and the queryset type
42+
for the provided queryset
43+
"""
44+
try:
45+
# Assume a queryset of nodes was provided
46+
_NodeModel = queryset.model
47+
_EdgeModel = queryset.model._meta.get_field("parents").through
48+
queryset_type = "nodes_queryset"
49+
except FieldDoesNotExist:
50+
try:
51+
# Assume a queryset of edges was provided
52+
_EdgeModel = queryset.model
53+
_NodeModel = queryset.model._meta.get_field("parent").related_model
54+
queryset_type = "edges_queryset"
55+
except FieldDoesNotExist:
56+
raise GraphModelsCannotBeParsedException
57+
return (_NodeModel, _EdgeModel, queryset_type)
58+
59+
60+
def model_to_dict(instance, fields=None, date_strf=None):
61+
"""
62+
Returns a dictionary of {field_name: field_value} for a given model instance
63+
e.g.: model_to_dict(myqueryset.first(), fields=["id",])
64+
65+
For DateTimeFields, a formatting string can be provided
3866
67+
Adapted from: https://ziwon.github.io/post/using_custom_model_to_dict_in_django/
68+
"""
3969

40-
def rawqueryset_to_dataframe(rawqueryset):
41-
"""Returns a pandas dataframe"""
42-
return pd.DataFrame(
43-
rawqueryset_to_values_list(rawqueryset), columns=list(rawqueryset.columns)
44-
)
70+
if not fields:
71+
raise IncorrectUsageException("fields list must be provided")
72+
73+
opts = instance._meta
74+
data = {}
75+
__fields = list(map(lambda a: a.split("__")[0], fields or []))
76+
77+
for f in chain(opts.concrete_fields, opts.private_fields, opts.many_to_many):
78+
is_editable = getattr(f, "editable", False)
79+
80+
if fields and f.name not in __fields:
81+
continue
82+
83+
if isinstance(f, DateTimeField):
84+
dt = f.value_from_object(instance)
85+
# Format based on format string provided, otherwise return a timestamp
86+
data[f.name] = dt.strftime(date_strf) if date_strf else dt.timestamp()
87+
88+
elif isinstance(f, ImageField):
89+
image = f.value_from_object(instance)
90+
data[f.name] = image.url if image else None
91+
92+
elif isinstance(f, FileField):
93+
file = f.value_from_object(instance)
94+
data[f.name] = file.url if file else None
95+
96+
elif isinstance(f, ManyToManyField):
97+
if instance.pk is None:
98+
data[f.name] = []
99+
else:
100+
qs = f.value_from_object(instance)
101+
if qs._result_cache is not None:
102+
data[f.name] = [item.pk for item in qs]
103+
else:
104+
try:
105+
m2m_field = list(
106+
filter(lambda a: f.name in a and a.find("__") != -1, fields)
107+
)[0]
108+
key = m2m_field[len(f.name) + 2 :]
109+
data[f.name] = list(qs.values_list(key, flat=True))
110+
except IndexError:
111+
data[f.name] = list(qs.values_list("pk", flat=True))
112+
113+
if isinstance(f, UUIDField):
114+
uuid = f.value_from_object(instance)
115+
data[f.name] = str(uuid) if uuid else None
116+
117+
# ToDo: Process other model fields
118+
119+
elif is_editable:
120+
data[f.name] = f.value_from_object(instance)
121+
122+
funcs = set(__fields) - set(list(data.keys()))
123+
for func in funcs:
124+
obj = getattr(instance, func)
125+
if inspect.ismethod(obj):
126+
data[func] = obj()
127+
else:
128+
data[func] = obj
129+
return data
130+
131+
132+
def edges_from_nodes_queryset(nodes_queryset):
133+
"""Given an Edge Model and a QuerySet or RawQuerySet of nodes,
134+
returns a queryset of the associated edges"""
135+
_NodeModel, _EdgeModel, queryset_type = get_queryset_characteristics(nodes_queryset)
136+
137+
if queryset_type == "nodes_queryset":
138+
return _filter_order(_EdgeModel.objects, ["parent", "child"], nodes_queryset)
139+
raise IncorrectQuerysetTypeException
140+
141+
142+
def nodes_from_edges_queryset(edges_queryset):
143+
"""Given a Node Model and a QuerySet or RawQuerySet of edges,
144+
returns a queryset of the associated nodes"""
145+
_NodeModel, _EdgeModel, queryset_type = get_queryset_characteristics(edges_queryset)
146+
147+
if queryset_type == "edges_queryset":
148+
149+
nodes_list = (
150+
_filter_order(
151+
_NodeModel.objects,
152+
[
153+
f"{_NodeModel.__name__}_child",
154+
],
155+
edges_queryset,
156+
)
157+
| _filter_order(
158+
_NodeModel.objects,
159+
[
160+
f"{_NodeModel.__name__}_parent",
161+
],
162+
edges_queryset,
163+
)
164+
).values_list("pk")
165+
166+
return _NodeModel.objects.filter(pk__in=nodes_list)
167+
raise IncorrectQuerysetTypeException
168+
169+
170+
def nx_from_queryset(
171+
queryset,
172+
graph_attributes_dict=None,
173+
node_attribute_fields_list=None,
174+
edge_attribute_fields_list=None,
175+
date_strf=None,
176+
):
177+
"""
178+
Provided a queryset of nodes or edges, returns a NetworkX graph
45179
180+
Optionally, the following can be supplied to add attributes to components of the generated graph:
181+
graph_attributes_dict: A dictionary of attributes to add to the graph itself
182+
node_attribute_fields_list: a list of strings of field names to be added to nodes
183+
edge_attribute_fields_list: a list of strings of field names to be added to edges
184+
"""
185+
_NodeModel, _EdgeModel, queryset_type = get_queryset_characteristics(queryset)
46186

47-
def edges_from_nodes_queryset(edge_model, nodes_queryset):
48-
"""Given an Edge Model and a QuerySet or RawQuerySet of nodes, returns a queryset of the associated edges"""
49-
return _filter_order(edge_model.objects, ["parent", "child"], nodes_queryset)
187+
if graph_attributes_dict is None:
188+
graph_attributes_dict = {}
50189

190+
graph = nx.Graph(**graph_attributes_dict)
51191

52-
def nodes_from_edges_queryset(node_model, edges_queryset):
53-
"""Given a Node Model and a QuerySet or RawQuerySet of edges, returns a queryset of the associated nodes"""
54-
nodes_list = (
55-
_filter_order(
56-
node_model.objects,
57-
[
58-
f"{node_model.__name__}_child",
59-
],
60-
edges_queryset,
61-
)
62-
| _filter_order(
63-
node_model.objects,
64-
[
65-
f"{node_model.__name__}_parent",
66-
],
67-
edges_queryset,
68-
)
69-
).values_list("pk")
192+
if queryset_type == "nodes_queryset":
193+
nodes_queryset = queryset
194+
edges_queryset = edges_from_nodes_queryset(nodes_queryset)
195+
else:
196+
edges_queryset = queryset
197+
nodes_queryset = nodes_from_edges_queryset(edges_queryset)
70198

71-
return node_model.objects.filter(pk__in=nodes_list)
199+
for node in nodes_queryset:
200+
if node_attribute_fields_list is not None:
201+
node_attribute_fields_dict = model_to_dict(
202+
node, fields=node_attribute_fields_list, date_strf=date_strf
203+
)
204+
else:
205+
node_attribute_fields_dict = {}
72206

207+
graph.add_node(node.pk, **node_attribute_fields_dict)
73208

74-
def nx_from_nodes_queryset(nodes_queryset):
75-
"""Provided a queryset of nodes, returns a NetworkX graph"""
76-
# ToDo: Implement
77-
pass
209+
for edge in edges_queryset:
210+
if edge_attribute_fields_list is not None:
211+
edge_attribute_fields_dict = model_to_dict(
212+
edge, fields=edge_attribute_fields_list, date_strf=date_strf
213+
)
214+
else:
215+
edge_attribute_fields_dict = {}
78216

217+
graph.add_edge(edge.parent.pk, edge.child.pk, **edge_attribute_fields_dict)
79218

80-
def nx_from_edges_queryset(edges_queryset, fields_array=None):
81-
"""Provided a queryset of edges, returns a NetworkX graph"""
82-
# ToDo: Implement
83-
graph = nx.Graph()
219+
return graph

setup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import os
44
from setuptools import setup
55

6-
version = '0.0.18'
6+
version = '0.0.19'
77

88
classifiers = [
99
"Development Status :: 3 - Alpha",

0 commit comments

Comments
 (0)