|
| 1 | +import ast |
1 | 2 | import csv |
2 | 3 | import io |
3 | 4 | import json |
|
11 | 12 | from dateutil.parser import parse as parse_datetime_str |
12 | 13 | from jinja2 import Template |
13 | 14 | from xhtml2pdf import pisa |
| 15 | +from pyvis.network import Network |
| 16 | +import networkx as nx |
14 | 17 |
|
| 18 | +from .checking import SUPPORTED_IDS |
15 | 19 | from .result import QueryStatus |
| 20 | +from .sites import MaigretDatabase |
16 | 21 | from .utils import is_country_tag, CaseConverter, enrich_link_str |
17 | 22 |
|
18 | 23 | SUPPORTED_JSON_REPORT_FORMATS = [ |
@@ -82,6 +87,121 @@ def save_json_report(filename: str, username: str, results: dict, report_type: s |
82 | 87 | generate_json_report(username, results, f, report_type=report_type) |
83 | 88 |
|
84 | 89 |
|
| 90 | +class MaigretGraph: |
| 91 | + other_params = {'size': 10, 'group': 3} |
| 92 | + site_params = {'size': 15, 'group': 2} |
| 93 | + username_params = {'size': 20, 'group': 1} |
| 94 | + |
| 95 | + def __init__(self, graph): |
| 96 | + self.G = graph |
| 97 | + |
| 98 | + def add_node(self, key, value): |
| 99 | + node_name = f'{key}: {value}' |
| 100 | + |
| 101 | + params = self.other_params |
| 102 | + if key in SUPPORTED_IDS: |
| 103 | + params = self.username_params |
| 104 | + elif value.startswith('http'): |
| 105 | + params = self.site_params |
| 106 | + |
| 107 | + self.G.add_node(node_name, title=node_name, **params) |
| 108 | + |
| 109 | + if value != value.lower(): |
| 110 | + normalized_node_name = self.add_node(key, value.lower()) |
| 111 | + self.link(node_name, normalized_node_name) |
| 112 | + |
| 113 | + return node_name |
| 114 | + |
| 115 | + def link(self, node1_name, node2_name): |
| 116 | + self.G.add_edge(node1_name, node2_name, weight=2) |
| 117 | + |
| 118 | + |
| 119 | +def save_graph_report(filename: str, username_results: list, db: MaigretDatabase): |
| 120 | + G = nx.Graph() |
| 121 | + graph = MaigretGraph(G) |
| 122 | + |
| 123 | + for username, id_type, results in username_results: |
| 124 | + username_node_name = graph.add_node(id_type, username) |
| 125 | + |
| 126 | + for website_name in results: |
| 127 | + dictionary = results[website_name] |
| 128 | + # TODO: fix no site data issue |
| 129 | + if not dictionary: |
| 130 | + continue |
| 131 | + |
| 132 | + if dictionary.get("is_similar"): |
| 133 | + continue |
| 134 | + |
| 135 | + status = dictionary.get("status") |
| 136 | + if not status: # FIXME: currently in case of timeout |
| 137 | + continue |
| 138 | + |
| 139 | + if dictionary["status"].status != QueryStatus.CLAIMED: |
| 140 | + continue |
| 141 | + |
| 142 | + site_fallback_name = dictionary.get('url_user', f'{website_name}: {username.lower()}') |
| 143 | + # site_node_name = dictionary.get('url_user', f'{website_name}: {username.lower()}') |
| 144 | + site_node_name = graph.add_node('site', site_fallback_name) |
| 145 | + graph.link(username_node_name, site_node_name) |
| 146 | + |
| 147 | + def process_ids(parent_node, ids): |
| 148 | + for k, v in ids.items(): |
| 149 | + if k.endswith('_count') or k.startswith('is_') or k.endswith('_at'): |
| 150 | + continue |
| 151 | + if k in 'image': |
| 152 | + continue |
| 153 | + |
| 154 | + v_data = v |
| 155 | + if v.startswith('['): |
| 156 | + try: |
| 157 | + v_data = ast.literal_eval(v) |
| 158 | + except Exception as e: |
| 159 | + logging.error(e) |
| 160 | + |
| 161 | + # value is a list |
| 162 | + if isinstance(v_data, list): |
| 163 | + list_node_name = graph.add_node(k, site_fallback_name) |
| 164 | + for vv in v_data: |
| 165 | + data_node_name = graph.add_node(vv, site_fallback_name) |
| 166 | + graph.link(list_node_name, data_node_name) |
| 167 | + |
| 168 | + add_ids = {a: b for b, a in db.extract_ids_from_url(vv).items()} |
| 169 | + if add_ids: |
| 170 | + process_ids(data_node_name, add_ids) |
| 171 | + else: |
| 172 | + # value is just a string |
| 173 | + # ids_data_name = f'{k}: {v}' |
| 174 | + # if ids_data_name == parent_node: |
| 175 | + # continue |
| 176 | + |
| 177 | + ids_data_name = graph.add_node(k, v) |
| 178 | + # G.add_node(ids_data_name, size=10, title=ids_data_name, group=3) |
| 179 | + graph.link(parent_node, ids_data_name) |
| 180 | + |
| 181 | + # check for username |
| 182 | + if 'username' in k or k in SUPPORTED_IDS: |
| 183 | + new_username_node_name = graph.add_node('username', v) |
| 184 | + graph.link(ids_data_name, new_username_node_name) |
| 185 | + |
| 186 | + add_ids = {k: v for v, k in db.extract_ids_from_url(v).items()} |
| 187 | + if add_ids: |
| 188 | + process_ids(ids_data_name, add_ids) |
| 189 | + |
| 190 | + if status.ids_data: |
| 191 | + process_ids(site_node_name, status.ids_data) |
| 192 | + |
| 193 | + nodes_to_remove = [] |
| 194 | + for node in G.nodes: |
| 195 | + if len(str(node)) > 100: |
| 196 | + nodes_to_remove.append(node) |
| 197 | + |
| 198 | + [G.remove_node(node) for node in nodes_to_remove] |
| 199 | + |
| 200 | + nt = Network(notebook=True, height="750px", width="100%") |
| 201 | + nt.from_nx(G) |
| 202 | + nt.show(filename) |
| 203 | + |
| 204 | + |
85 | 205 | def get_plaintext_report(context: dict) -> str: |
86 | 206 | output = (context['brief'] + " ").replace('. ', '.\n') |
87 | 207 | interests = list(map(lambda x: x[0], context.get('interests_tuple_list', []))) |
|
0 commit comments