refactor

app.py

@@ -1,29 +1,11 @@
 import streamlit as st
 import asyncio
-import websockets
-import json
-import tempfile
-import os
-from typing import Dict, Any, Optional, List, Literal, TypedDict, cast, Set, Tuple
-from streamlit_agraph import agraph, Node, Edge, Config
-import pyvis.network as net
+from typing import Dict, Any, Optional
+from streamlit_agraph import agraph, Config
 
-GGRAPHER_URI = "wss://ggrphr.davidalber.net"
-
-class StartNodeRequest(TypedDict):
-    recordId: int
-    getAdvisors: bool
-    getDescendants: bool
-
-class RequestPayload(TypedDict):
-    kind: Literal["build-graph"]
-    options: Dict[Literal["reportingCallback"], bool]
-    startNodes: List[StartNodeRequest]
-
-class ProgressCallback(TypedDict):
-    queued: int
-    fetching: int
-    done: int
+from src.network import make_payload, get_graph
+from src.graph import build_tree_structure, create_hierarchical_view, tree_to_dot
+from src.viz import create_pyvis_network
 
 def get_id_from_input(val: str) -> Optional[int]:
     try:
@@ -31,265 +13,16 @@ def get_id_from_input(val: str) -> Optional[int]:
     except Exception:
         return None
 
-def make_payload(record_id: int) -> RequestPayload:
-    return {
-        "kind": "build-graph",
-        "options": {"reportingCallback": True},
-        "startNodes": [{
-            "recordId": record_id,
-            "getAdvisors": True,
-            "getDescendants": False,
-        }],
-    }
-
-async def get_graph(payload: RequestPayload, progress_cb=None) -> Dict[str, Any]:
-    def intify_record_keys(d: Dict[Any, Any]) -> Dict[Any, Any]:
-        if "nodes" in d:
-            ret = {k: v for k, v in d.items() if k != "nodes"}
-            ret["nodes"] = {int(k): v for k, v in d["nodes"].items()}
-            return ret
-        return d
-
-    async with websockets.connect(  # type: ignore[attr-defined]
-        GGRAPHER_URI,
-    ) as ws:
-        await ws.send(json.dumps(payload))
-        while True:
-            response_json = await ws.recv()
-            response = json.loads(response_json, object_hook=intify_record_keys)
-            response_payload = response.get("payload")
-            if response["kind"] == "graph":
-                return cast(Dict[str, Any], response_payload)
-            elif response["kind"] == "progress" and progress_cb:
-                progress = cast(ProgressCallback, response_payload)
-                progress_cb(progress)
-            else:
-                continue
-
-def build_tree_structure(graph: Dict[str, Any], root_id: int) -> Dict[int, Dict[str, Any]]:
-    """Build a hierarchical tree structure from the graph data."""
-    nodes = graph.get("nodes", {})
-    tree = {}
-    
-    # calculate depth/generation for each node
-    def calculate_depths(node_id: int, visited: Set[int], depth: int = 0) -> Dict[int, int]:
-        if node_id in visited:
-            return {}
-        
-        visited.add(node_id)
-        depths = {node_id: depth}
-        
-        node = nodes.get(node_id, {})
-        for advisor_id in node.get("advisors", []):
-            if advisor_id in nodes:
-                advisor_depths = calculate_depths(advisor_id, visited, depth + 1)
-                depths.update(advisor_depths)
-        
-        return depths
-    
-    depths = calculate_depths(root_id, set())
-    
-    # build tree structure with depth info
-    for node_id, node in nodes.items():
-        if node_id in depths:
-            tree[node_id] = {
-                **node,
-                "depth": depths[node_id],
-                "children": [],
-                "advisors": node.get("advisors", [])
-            }
-    
-    # establish parent-child relationships
-    for node_id, node in tree.items():
-        for advisor_id in node["advisors"]:
-            if advisor_id in tree:
-                tree[advisor_id]["children"].append(node_id)
-    
-    return tree
-
-def create_hierarchical_view(graph: Dict[str, Any], root_id: int, max_depth: int = None) -> Tuple[List[Node], List[Edge]]:
-    """Create nodes and edges for hierarchical tree view."""
+def display_tree_summary(graph: Dict[str, Any], root_id: int) -> None:
     tree = build_tree_structure(graph, root_id)
-    nodes_list = []
-    edges_list = []
-    
-    # color scheme by generation
-    colors = ["#ff6b6b", "#4ecdc4", "#45b7d1", "#96ceb4", "#feca57", "#ff9ff3", "#54a0ff"]
-    
-    # group nodes by depth for positioning
-    nodes_by_depth = {}
-    for node_id, node in tree.items():
-        if max_depth is None or node["depth"] <= max_depth:
-            depth = node["depth"]
-            if depth not in nodes_by_depth:
-                nodes_by_depth[depth] = []
-            nodes_by_depth[depth].append((node_id, node))
-    
-    # sort nodes within each depth by year (oldest first, recent last for lower position)
-    for depth in nodes_by_depth:
-        nodes_by_depth[depth].sort(key=lambda x: x[1].get('year') or 1400)
-    
-    # create nodes with positioning hints
-    for depth in sorted(nodes_by_depth.keys()):
-        depth_nodes = nodes_by_depth[depth]
-        
-        for i, (node_id, node) in enumerate(depth_nodes):
-            color = colors[depth % len(colors)]
-            
-            name = node.get("name", str(node_id))
-            year_str = f" ({node.get('year')})" if node.get('year') is not None else ""
-            label = f"{name}{year_str}"
-            
-            # calculate positioning based on year within generation
-            year = node.get('year') or 1500  # default to old year if missing
-            base_y = depth * 300  # spacing between generations
-            
-            # position nodes with year-based offset within generation
-            # more recent years get higher y values (appear lower on screen)
-            year_offset = (year - 1400) * 0.2  # scale factor for year spacing
-            x_pos = i * 180 + (depth * 20)  # slight x offset per depth to avoid overlap
-            y_pos = base_y + year_offset
-            
-            # create streamlit-agraph node with positioning
-            ag_node = Node(
-                id=str(node_id),
-                label=label,
-                size=25 if node_id == root_id else 20,
-                color=color,
-                title=f"Name: {name}\nYear: {node.get('year', 'N/A')}\nInstitution: {node.get('institution', 'N/A')}",
-                x=x_pos,
-                y=y_pos,
-                font={"color": "white", "size": 12}
-            )
-            nodes_list.append(ag_node)
-            
-            # create edges to advisors
-            for advisor_id in node["advisors"]:
-                if advisor_id in tree and (max_depth is None or tree[advisor_id]["depth"] <= max_depth):
-                    edge = Edge(
-                        source=str(advisor_id),
-                        target=str(node_id),
-                        color="#666666"
-                    )
-                    edges_list.append(edge)
-    
-    return nodes_list, edges_list
-
-def create_pyvis_network(graph: Dict[str, Any], root_id: int) -> str:
-    """Create an interactive network using pyvis."""
-    nodes = graph.get("nodes", {})
-    
-    # create network
-    nt = net.Network(
-        height="600px",
-        width="100%",
-        bgcolor="#ffffff",
-        font_color="black",
-        directed=True
-    )
-    
-    # configure physics
-    nt.set_options("""
-    var options = {
-      "physics": {
-        "enabled": true,
-        "hierarchicalRepulsion": {
-          "centralGravity": 0.3,
-          "springLength": 120,
-          "springConstant": 0.01,
-          "nodeDistance": 200,
-          "damping": 0.09
-        },
-        "solver": "hierarchicalRepulsion"
-      },
-      "layout": {
-        "hierarchical": {
-          "enabled": true,
-          "direction": "UD",
-          "sortMethod": "directed"
-        }
-      }
-    }
-    """)
-    
-    # calculate depths for coloring
-    tree = build_tree_structure(graph, root_id)
-    colors = ["#ff6b6b", "#4ecdc4", "#45b7d1", "#96ceb4", "#feca57", "#ff9ff3", "#54a0ff"]
-    
-    # add nodes
-    for node_id, node in nodes.items():
-        name = node.get("name", str(node_id))
-        year_str = f" ({node.get('year')})" if node.get('year') is not None else ""
-        label = f"{name}{year_str}"
-        
-        depth = tree.get(node_id, {}).get("depth", 0)
-        color = colors[depth % len(colors)]
-        
-        nt.add_node(
-            node_id,
-            label=label,
-            title=f"Name: {name}\nYear: {node.get('year', 'N/A')}\nInstitution: {node.get('institution', 'N/A')}",
-            color=color,
-            size=30 if node_id == root_id else 20
-        )
-    
-    # add edges
-    for node_id, node in nodes.items():
-        for advisor_id in node.get("advisors", []):
-            if advisor_id in nodes:
-                nt.add_edge(advisor_id, node_id)
-    
-    # save to temp file
-    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".html")
-    nt.save_graph(temp_file.name)
-    
-    with open(temp_file.name, 'r') as f:
-        html_content = f.read()
-    
-    os.unlink(temp_file.name)
-    return html_content
-
-def tree_to_dot(graph: Dict[str, Any]) -> str:
-    """Original graphviz DOT format (kept as fallback)."""
-    nodes = graph.get("nodes", {})
-    lines = [
-        "digraph G {",
-        "  rankdir=TB;",
-        '  node [shape=box, style="rounded,filled", fillcolor=lightyellow];',
-        '  edge [arrowhead=vee];'
-    ]
-    
-    for node_id, node in nodes.items():
-        name = node.get("name", str(node_id))
-        year_str = f" ({node.get('year')})" if node.get('year') is not None else " (Year Unknown)"
-        label = f"{name}{year_str}"
-        tooltip = f"ID: {node_id}\\nName: {name}\\nYear: {node.get('year', 'N/A')}\\nInstitution: {node.get('institution', 'N/A')}"
-        lines.append(f'  "{node_id}" [label="{label}", tooltip="{tooltip}"];')
-    
-    for node_id, node in nodes.items():
-        for adv_id in node.get("advisors", []):
-            if adv_id in nodes:
-                lines.append(f'  "{adv_id}" -> "{node_id}";')
-    
-    lines.append("}")
-    return "\n".join(lines)
-
-def display_tree_summary(graph: Dict[str, Any], root_id: int):
-    """Display summary statistics about the tree."""
-    tree = build_tree_structure(graph, root_id)
-    
     if not tree:
         return
-    
     max_depth = max(node["depth"] for node in tree.values()) if tree else 0
     total_nodes = len(tree)
-    
-    # count nodes by generation
-    depth_counts = {}
+    depth_counts: dict[int, int] = {}
     for node in tree.values():
         depth = node["depth"]
         depth_counts[depth] = depth_counts.get(depth, 0) + 1
-    
     col1, col2, col3 = st.columns(3)
     with col1:
         st.metric("Total Mathematicians", total_nodes)
@@ -298,8 +31,6 @@ def display_tree_summary(graph: Dict[str, Any], root_id: int):
     with col3:
         root_name = tree.get(root_id, {}).get("name", "Unknown")
         st.metric("Root", root_name)
-    
-
 def main():
     st.title("Math Genealogy Ancestor Tree")
     st.write("Interactive visualization of academic advisor relationships from the Mathematics Genealogy Project")

pyproject.toml

@@ -9,13 +9,16 @@ packages = [{ include = "src" }]
 [tool.poetry.dependencies]
 python = ">=3.9,<3.9.7 || >3.9.7,<4.0"
 streamlit = "^1.35.0"
-aiohttp = "^3.12.0"
-beautifulsoup4 = "^4.13.4"
 streamlit-agraph = "^0.0.45"
 pyvis = "^0.3.2"
+websockets = "^12.0"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.2.0"
+black = "^24.4.2"
+isort = "^5.13.2"
+flake8 = "^7.0.0"
+mypy = "^1.10.0"
 
 [build-system]
 requires = ["poetry-core"]

src/graph.py

@@ -0,0 +1,111 @@
+from typing import Dict, Any, Set, Tuple, List
+from streamlit_agraph import Node, Edge
+
+def build_tree_structure(graph: Dict[str, Any], root_id: int) -> Dict[int, Dict[str, Any]]:
+    nodes = graph.get("nodes", {})
+    tree = {}
+
+    def calculate_depths(node_id: int, visited: Set[int], depth: int = 0) -> Dict[int, int]:
+        if node_id in visited:
+            return {}
+        visited.add(node_id)
+        depths = {node_id: depth}
+        node = nodes.get(node_id, {})
+        for advisor_id in node.get("advisors", []):
+            if advisor_id in nodes:
+                advisor_depths = calculate_depths(advisor_id, visited, depth + 1)
+                depths.update(advisor_depths)
+        return depths
+
+    depths = calculate_depths(root_id, set())
+
+    for node_id, node in nodes.items():
+        if node_id in depths:
+            tree[node_id] = {
+                **node,
+                "depth": depths[node_id],
+                "children": [],
+                "advisors": node.get("advisors", [])
+            }
+
+    for node_id, node in tree.items():
+        for advisor_id in node["advisors"]:
+            if advisor_id in tree:
+                tree[advisor_id]["children"].append(node_id)
+
+    return tree
+
+def create_hierarchical_view(
+    graph: Dict[str, Any],
+    root_id: int,
+    max_depth: int = None
+) -> Tuple[List[Node], List[Edge]]:
+    tree = build_tree_structure(graph, root_id)
+    nodes_list = []
+    edges_list = []
+
+    colors = ["#ff6b6b", "#4ecdc4", "#45b7d1", "#96ceb4", "#feca57", "#ff9ff3", "#54a0ff"]
+    nodes_by_depth = {}
+    for node_id, node in tree.items():
+        if max_depth is None or node["depth"] <= max_depth:
+            depth = node["depth"]
+            if depth not in nodes_by_depth:
+                nodes_by_depth[depth] = []
+            nodes_by_depth[depth].append((node_id, node))
+
+    for depth in nodes_by_depth:
+        nodes_by_depth[depth].sort(key=lambda x: x[1].get('year') or 1400)
+
+    for depth in sorted(nodes_by_depth.keys()):
+        depth_nodes = nodes_by_depth[depth]
+        for i, (node_id, node) in enumerate(depth_nodes):
+            color = colors[depth % len(colors)]
+            name = node.get("name", str(node_id))
+            year_str = f" ({node.get('year')})" if node.get('year') is not None else ""
+            label = f"{name}{year_str}"
+            year = node.get('year') or 1500
+            base_y = depth * 300
+            year_offset = (year - 1400) * 0.2
+            x_pos = i * 180 + (depth * 20)
+            y_pos = base_y + year_offset
+            ag_node = Node(
+                id=str(node_id),
+                label=label,
+                size=25 if node_id == root_id else 20,
+                color=color,
+                title=f"Name: {name}\nYear: {node.get('year', 'N/A')}\nInstitution: {node.get('institution', 'N/A')}",
+                x=x_pos,
+                y=y_pos,
+                font={"color": "white", "size": 12}
+            )
+            nodes_list.append(ag_node)
+            for advisor_id in node["advisors"]:
+                if advisor_id in tree and (max_depth is None or tree[advisor_id]["depth"] <= max_depth):
+                    edge = Edge(
+                        source=str(advisor_id),
+                        target=str(node_id),
+                        color="#666666"
+                    )
+                    edges_list.append(edge)
+    return nodes_list, edges_list
+
+def tree_to_dot(graph: Dict[str, Any]) -> str:
+    nodes = graph.get("nodes", {})
+    lines = [
+        "digraph G {",
+        "  rankdir=TB;",
+        '  node [shape=box, style="rounded,filled", fillcolor=lightyellow];',
+        '  edge [arrowhead=vee];'
+    ]
+    for node_id, node in nodes.items():
+        name = node.get("name", str(node_id))
+        year_str = f" ({node.get('year')})" if node.get('year') is not None else " (Year Unknown)"
+        label = f"{name}{year_str}"
+        tooltip = f"ID: {node_id}\\nName: {name}\\nYear: {node.get('year', 'N/A')}\\nInstitution: {node.get('institution', 'N/A')}"
+        lines.append(f'  "{node_id}" [label="{label}", tooltip="{tooltip}"];')
+    for node_id, node in nodes.items():
+        for adv_id in node.get("advisors", []):
+            if adv_id in nodes:
+                lines.append(f'  "{adv_id}" -> "{node_id}";')
+    lines.append("}")
+    return "\n".join(lines)

src/network.py

@@ -0,0 +1,57 @@
+import asyncio
+import websockets
+import json
+from typing import Dict, Any, Callable, Optional, TypedDict, Literal, List, cast
+
+GGRAPHER_URI = "wss://ggrphr.davidalber.net"
+
+class StartNodeRequest(TypedDict):
+    recordId: int
+    getAdvisors: bool
+    getDescendants: bool
+
+class RequestPayload(TypedDict):
+    kind: Literal["build-graph"]
+    options: Dict[Literal["reportingCallback"], bool]
+    startNodes: List[StartNodeRequest]
+
+class ProgressCallback(TypedDict):
+    queued: int
+    fetching: int
+    done: int
+
+def make_payload(record_id: int) -> RequestPayload:
+    return {
+        "kind": "build-graph",
+        "options": {"reportingCallback": True},
+        "startNodes": [{
+            "recordId": record_id,
+            "getAdvisors": True,
+            "getDescendants": False,
+        }],
+    }
+
+async def get_graph(
+    payload: RequestPayload,
+    progress_cb: Optional[Callable[[ProgressCallback], None]] = None
+) -> Dict[str, Any]:
+    def intify_record_keys(d: Dict[Any, Any]) -> Dict[Any, Any]:
+        if "nodes" in d:
+            ret = {k: v for k, v in d.items() if k != "nodes"}
+            ret["nodes"] = {int(k): v for k, v in d["nodes"].items()}
+            return ret
+        return d
+
+    async with websockets.connect(GGRAPHER_URI) as ws:
+        await ws.send(json.dumps(payload))
+        while True:
+            response_json = await ws.recv()
+            response = json.loads(response_json, object_hook=intify_record_keys)
+            response_payload = response.get("payload")
+            if response["kind"] == "graph":
+                return cast(Dict[str, Any], response_payload)
+            elif response["kind"] == "progress" and progress_cb:
+                progress = cast(ProgressCallback, response_payload)
+                progress_cb(progress)
+            else:
+                continue

src/viz.py

@@ -0,0 +1,62 @@
+import tempfile
+import os
+from typing import Dict, Any
+import pyvis.network as net
+from .graph import build_tree_structure
+
+def create_pyvis_network(graph: Dict[str, Any], root_id: int) -> str:
+    nodes = graph.get("nodes", {})
+    nt = net.Network(
+        height="600px",
+        width="100%",
+        bgcolor="#ffffff",
+        font_color="black",
+        directed=True
+    )
+    nt.set_options("""
+    var options = {
+      "physics": {
+        "enabled": true,
+        "hierarchicalRepulsion": {
+          "centralGravity": 0.3,
+          "springLength": 120,
+          "springConstant": 0.01,
+          "nodeDistance": 200,
+          "damping": 0.09
+        },
+        "solver": "hierarchicalRepulsion"
+      },
+      "layout": {
+        "hierarchical": {
+          "enabled": true,
+          "direction": "UD",
+          "sortMethod": "directed"
+        }
+      }
+    }
+    """)
+    tree = build_tree_structure(graph, root_id)
+    colors = ["#ff6b6b", "#4ecdc4", "#45b7d1", "#96ceb4", "#feca57", "#ff9ff3", "#54a0ff"]
+    for node_id, node in nodes.items():
+        name = node.get("name", str(node_id))
+        year_str = f" ({node.get('year')})" if node.get('year') is not None else ""
+        label = f"{name}{year_str}"
+        depth = tree.get(node_id, {}).get("depth", 0)
+        color = colors[depth % len(colors)]
+        nt.add_node(
+            node_id,
+            label=label,
+            title=f"Name: {name}\nYear: {node.get('year', 'N/A')}\nInstitution: {node.get('institution', 'N/A')}",
+            color=color,
+            size=30 if node_id == root_id else 20
+        )
+    for node_id, node in nodes.items():
+        for advisor_id in node.get("advisors", []):
+            if advisor_id in nodes:
+                nt.add_edge(advisor_id, node_id)
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".html")
+    nt.save_graph(temp_file.name)
+    with open(temp_file.name, 'r') as f:
+        html_content = f.read()
+    os.unlink(temp_file.name)
+    return html_content