r/greens_report
1
0
Fork 0
forked from security_vision/semantic_graph
Code Issues Pull requests Projects Releases Wiki Activity

Option to generate nodes and edges CSV files

Browse source
This commit is contained in:
Ruben van de Ven 2023-01-23 21:55:28 +01:00
parent 3c0184fe97
commit d457c8eea0
1 changed files with 20 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
wiki_relations.py
View file

@ -4,6 +4,7 @@ import requests
import argparse import argparse
import datetime import datetime
import tqdm import tqdm
import csv
@ -12,11 +13,12 @@ logger = logging.getLogger('wiki')
default_categories = [ default_categories = [
# 'Person', # 'Person',
'Institution', 'Institution',
'Technology', 'Products',
'Deployments', 'Deployments',
'Dataset', 'Dataset',
'City', 'City',
# 'Country',# for deployments without city we should configure Geolocation #'Country',# for deployments without city we should configure Geolocation
'Technology Type',
] ]
parser = argparse.ArgumentParser(description='Turn wiki into nodes & links, usable by d3-force.') parser = argparse.ArgumentParser(description='Turn wiki into nodes & links, usable by d3-force.')
@ -28,6 +30,8 @@ parser.add_argument('--output', default="semantic_data.json",
help='Output JSON file') help='Output JSON file')
parser.add_argument('--credentials', default="no_credentials.json", parser.add_argument('--credentials', default="no_credentials.json",
help="JSON file containing the Bot's credentials") help="JSON file containing the Bot's credentials")
parser.add_argument('--generate-csv', action='store_true',
help="generate edge.csv & nodes.csv")
args = parser.parse_args() args = parser.parse_args()
@ -257,3 +261,17 @@ if __name__ == "__main__":
with open(args.output, 'w') as fp: with open(args.output, 'w') as fp:
json.dump(collection, fp) json.dump(collection, fp)
if args.generate_csv:
with open('nodes.csv', 'w') as csvfile:
all_keys = set().union(*(d.keys() for d in collection['nodes']))
# all_keys = ['@id']
dict_writer = csv.DictWriter(csvfile, fieldnames=all_keys, extrasaction='ignore', restval='')
dict_writer.writeheader()
dict_writer.writerows(collection['nodes'])
with open('edges.csv', 'w') as csvfile:
all_keys = set().union(*(d.keys() for d in collection['links']))
dict_writer = csv.DictWriter(csvfile, fieldnames=all_keys, extrasaction='ignore', restval='')
dict_writer.writeheader()
dict_writer.writerows(collection['links'])