156 lines
4.2 KiB
Nim
156 lines
4.2 KiB
Nim
# vim: set et sta sw=4 ts=4 :
|
|
#
|
|
# Find links between two actors, via other actors
|
|
# they've worked with across movies.
|
|
#
|
|
# Outputs to screen, and generates a dot file with
|
|
# results for graphviz visualization.
|
|
#
|
|
# Compile:
|
|
# % nim c -d:release imdb_find_actor_path.nim
|
|
#
|
|
# Run the "imdb_import" utility before using this!
|
|
#
|
|
|
|
import
|
|
std/math,
|
|
std/os,
|
|
std/sequtils,
|
|
std/strformat
|
|
import kuzu
|
|
|
|
const DB = "imdb.kz"
|
|
const DOT = "imdb-results.dot"
|
|
|
|
if not DB.fileExists:
|
|
echo """Cowardly refusing to run without an imdb database.
|
|
(see: imdb_import in this directory.)"""
|
|
quit 1
|
|
|
|
if paramCount() < 2:
|
|
echo "I require 2 actor names, in quotes."
|
|
quit 1
|
|
|
|
var stmt: KuzuPreparedStatement
|
|
var res: KuzuQueryResult
|
|
|
|
var fromActor = paramStr(1)
|
|
var toActor = paramStr(2)
|
|
|
|
var db = newKuzuDatabase( DB, kuzuConfig(read_only=true) )
|
|
var conn = db.connect
|
|
|
|
echo "Database opened: ", db.path
|
|
|
|
stmt = conn.prepare( "MATCH (a:Actor {name:$actor}) RETURN count(a)" )
|
|
res = stmt.execute( (actor: fromActor) )
|
|
if res.getNext[0].toInt64 < 1:
|
|
echo "Couldn't find actor ", &"\"{fromActor}\"."
|
|
quit 1
|
|
|
|
stmt = conn.prepare( "MATCH (a:Actor {name:$actor}) RETURN count(a)" )
|
|
res = stmt.execute( (actor: toActor) )
|
|
if res.getNext[0].toInt64 < 1:
|
|
echo "Couldn't find actor ", &"\"{toActor}\"."
|
|
quit 1
|
|
|
|
stmt = conn.prepare """
|
|
MATCH path = (a:Actor {name:$fromActor})-[r:ActedIn* ALL SHORTEST]-(b:Actor {name:$toActor})
|
|
RETURN DISTINCT nodes(path) AS nodes, length(path) AS hops
|
|
ORDER BY hops
|
|
"""
|
|
stdout.write &"Finding paths from {fromActor} to {toActor}... "
|
|
stdout.flushFile
|
|
res = stmt.execute( (fromActor: fromActor, toActor: toActor) )
|
|
|
|
if res.num_tuples == 0:
|
|
echo "Unable to find any paths!"
|
|
quit 1
|
|
|
|
|
|
let dotFile = DOT.open( fmWrite )
|
|
dotFile.write "strict digraph {\n"
|
|
dotfile.write &"""
|
|
graph[
|
|
rankdir=LR,
|
|
size="8.5,11",
|
|
margin=0.4,
|
|
label="Tracing from \"{fromActor}\" to \"{toActor}\""
|
|
];
|
|
node [
|
|
fontname=Arial
|
|
];""", "\n\n"
|
|
|
|
var fastestPath = res.getNext[1].toInt64
|
|
echo &"{round(res.execution_time / 1000, 2)} seconds, ",
|
|
"fastest path in ", fastestPath, " hop(s).\n"
|
|
res.rewind
|
|
|
|
# 1st pass, get all nodes for per-node styles
|
|
#
|
|
var nodes: seq[ tuple[kind: char, label: string] ] = @[]
|
|
for row in res.items:
|
|
for rawNode in row[0].toList:
|
|
var node = rawNode.toNode
|
|
var kind = $node["_LABEL"]
|
|
case kind:
|
|
of "Actor":
|
|
nodes.add( (kind: 'a', label: $node["name"]) )
|
|
of "Movie":
|
|
nodes.add( (kind: 'm', label: $node["title"]) )
|
|
else:
|
|
discard
|
|
|
|
# Pre-define unique node styles
|
|
#
|
|
for node in nodes.deDuplicate():
|
|
case node.kind:
|
|
of 'a':
|
|
dotFile.write( &"\"{node.label}\" " )
|
|
if node.label == fromActor or node.label == toActor:
|
|
dotFile.write """[shape=box, style=rounded, penwidth=3.0];""", "\n"
|
|
else:
|
|
dotFile.write """[shape=box, style=rounded];""", "\n"
|
|
of 'm':
|
|
dotFile.write( &"\"{node.label}\" " )
|
|
dotFile.write """[shape=note];""", "\n"
|
|
else:
|
|
discard
|
|
|
|
# 2nd pass: emit relations to stdout and create dot links
|
|
#
|
|
dotFile.write "\n\n"
|
|
if res.num_tuples > 1:
|
|
echo res.num_tuples, " tied ", fastestPath, " hop paths:"
|
|
for row in res.items:
|
|
var output: string
|
|
var pathLen = row[1].toInt64
|
|
var pathStep = 0
|
|
|
|
for rawNode in row[0].toList:
|
|
pathStep += 1
|
|
var node = rawNode.toNode
|
|
if $node["_LABEL"] == "Actor":
|
|
output.add $node["name"]
|
|
dotFile.write &""""{$node["name"]}""""
|
|
if pathStep == 1:
|
|
output.add " was in "
|
|
dotFile.write " -> "
|
|
elif pathStep < pathLen:
|
|
output.add " who was in "
|
|
dotFile.write " -> "
|
|
dotFile.write "\n"
|
|
else:
|
|
dotFile.write ";\n"
|
|
elif $node["_LABEL"] == "Movie":
|
|
output.add &""""{$node["title"]}""""
|
|
output.add " with "
|
|
dotFile.write &""""{$node["title"]}""""
|
|
dotFile.write " -> \n"
|
|
echo &"{output}."
|
|
|
|
dotFile.write "}\n"
|
|
dotFile.close
|
|
echo "\n\nYou can run 'dot -Tpdf < imdb-results.dot > imdb-results.pdf' if you have graphviz installed."
|
|
|
|
|