Added option to choose directed or undirected graph export

OmenApps · Nov 25, 2020 · 1d44584 · 1d44584
1 parent f83d7b4
commit 1d44584
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/django_postgresql_dag/transformations.py b/django_postgresql_dag/transformations.py
@@ -194,6 +194,7 @@ def nx_from_queryset(
     node_attribute_fields_list=None,
     edge_attribute_fields_list=None,
     date_strf=None,
+    digraph=False,
 ):
     """
     Provided a queryset of nodes or edges, returns a NetworkX graph
@@ -208,7 +209,10 @@ def nx_from_queryset(
     if graph_attributes_dict is None:
         graph_attributes_dict = {}
 
-    graph = nx.Graph(**graph_attributes_dict)
+    if not digraph:
+        graph = nx.Graph(**graph_attributes_dict)
+    else:
+        graph = nx.DiGraph(**graph_attributes_dict)
 
     if queryset_type == "nodes_queryset":
         nodes_queryset = queryset

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 import os
 from setuptools import setup
 
-version = '0.1.5'
+version = '0.1.6'
 
 classifiers = [
     "Development Status :: 3 - Alpha",