Support string labels with CSV entity labels

grandlite/__init__.py

@@ -20,6 +20,7 @@
  "jsonl": lambda x: x.to_json(sys.stdout, orient="records", lines=True),
 }
 
+
 def _guess_delimiter(first_n_lines: list[str]) -> str:
  """
  Guess the delimiter of a CSV file from the first few lines.
@@ -38,10 +39,14 @@ def _guess_delimiter(first_n_lines: list[str]) -> str:
  # Appears at all?
  if not any(delimiter in line for line in first_n_lines):
  continue
- if all(line.count(delimiter) == first_n_lines[0].count(delimiter) for line in first_n_lines):
+ if all(
+ line.count(delimiter) == first_n_lines[0].count(delimiter)
+ for line in first_n_lines
+ ):
  return delimiter
  raise ValueError("Could not guess delimiter.")
 
+
 def read_headered_edgelist(filename: str) -> nx.Graph:
  """
  Read a graph from a headered edgelist file.
@@ -62,6 +67,7 @@ def read_headered_edgelist(filename: str) -> nx.Graph:
  src_col = match.group(1)
  tgt_col = match.group(2)
  filepath = match.group(3)
+
  # The file has a header row.
  # Use the CSV reader to read the file:
  def _without_srctgt(row):
@@ -108,7 +114,6 @@ def _infer_graph_filetype_from_contents(filename: str) -> str:
  if match is not None:
  return "opencypher"
 
-
  raise NotImplementedError("Cannot infer graph file type from contents.")
 
 
@@ -128,6 +133,21 @@ def read_opencypher(paths: str) -> nx.Graph:
  return opencypher_buffers_to_graph(vertex_paths, edge_paths)
 
 
+def parse_labels_attribute(labels_str: str) -> set:
+ """
+ Parse a CSV string of labels into a set.
+
+ Arguments:
+ labels_str: The CSV string of labels.
+
+ Returns:
+ A set of labels.
+ """
+ if labels_str:
+ return set(labels_str.split(","))
+ return set()
+
+
 def detect_and_load_graph(graph_uri: str) -> nx.Graph:
  """
  Read a graph from its URI and return a NetworkX.Graph-compatible API.
@@ -184,6 +204,15 @@ def detect_and_load_graph(graph_uri: str) -> nx.Graph:
  raise ValueError(f"Unknown graph file type for file '{graph_path}'.")
 
  host_graph = readers[graph_type](graph_path)
+ # Parse __labels__ attributes as CSV and convert to set
+ for nid, node_attrs in host_graph.nodes(data=True):
+ if "__labels__" in node_attrs:
+ node_attrs["__labels__"] = parse_labels_attribute(node_attrs["__labels__"])
+
+ for u, v, edge_attrs in host_graph.edges(data=True):
+ if "__labels__" in edge_attrs:
+ edge_attrs["__labels__"] = parse_labels_attribute(edge_attrs["__labels__"])
+
  return host_graph
 
 
@@ -312,6 +341,7 @@ def cli():
  elif args.output in ["json", "jsonl"]:
  # Write valid JSON to stdout
  import json
+
  json.dump(results, sys.stdout)
  else:
  writer = csv.DictWriter(sys.stdout, fieldnames=results.keys())