Skip to content

Commit b6e324c

Browse files
swilly22swilly22
committed
return entire entities
1 parent 0388b1b commit b6e324c

File tree

1 file changed

+84
-78
lines changed

1 file changed

+84
-78
lines changed

‎code_graph/graph.py

Lines changed: 84 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -241,6 +241,22 @@ def get_neighbors(self, node_id: int, rel: Optional[str] = None, lbl: Optional[s
241241
return {'nodes': [], 'edges': []}
242242

243243

244+
def _class_from_node(self, n: Node) -> Class:
245+
"""
246+
Create a Class from a graph node
247+
"""
248+
249+
doc = n.properties.get('doc')
250+
name = n.properties.get('name')
251+
path = n.properties.get('path')
252+
src_end = n.properties.get('src_end')
253+
src_start = n.properties.get('src_start')
254+
255+
c = Class(path, name, doc, src_start, src_end)
256+
c.id = n.id
257+
258+
return c
259+
244260
def add_class(self, c: Class) -> None:
245261
"""
246262
Adds a class node to the graph database.
@@ -252,7 +268,7 @@ def add_class(self, c: Class) -> None:
252268
q = """MERGE (c:Class:Searchable {name: $name, path: $path, src_start: $src_start,
253269
src_end: $src_end})
254270
SET c.doc = $doc
255-
RETURN ID(c)"""
271+
RETURN c"""
256272

257273
params = {
258274
'doc': c.doc,
@@ -262,24 +278,9 @@ def add_class(self, c: Class) -> None:
262278
'src_end': c.src_end,
263279
}
264280

265-
res = self._query(q, params)
266-
c.id = res.result_set[0][0]
267-
268-
def _class_from_node(self, n: Node) -> Class:
269-
"""
270-
Create a Class from a graph node
271-
"""
272-
273-
doc = n.properties.get('doc')
274-
name = n.properties.get('name')
275-
path = n.properties.get('path')
276-
src_end = n.properties.get('src_end')
277-
src_start = n.properties.get('src_start')
278-
279-
c = Class(path, name, doc, src_start, src_end)
280-
c.id = n.id
281-
282-
return c
281+
res = self._query(q, params)
282+
node = res.result_set[0][0]
283+
c.id = node.id
283284

284285
def get_class_by_name(self, class_name: str) -> Optional[Class]:
285286
q = "MATCH (c:Class) WHERE c.name = $name RETURN c LIMIT 1"
@@ -303,6 +304,30 @@ def get_class(self, class_id: int) -> Optional[Class]:
303304
c = res.result_set[0][0]
304305
return self._class_from_node(c)
305306

307+
def _function_from_node(self, n: Node) -> Function:
308+
"""
309+
Create a Function from a graph node
310+
"""
311+
312+
src = n.properties.get('src')
313+
doc = n.properties.get('doc')
314+
path = n.properties.get('path')
315+
name = n.properties.get('name')
316+
args = n.properties.get('args')
317+
src_end = n.properties.get('src_end')
318+
ret_type = n.properties.get('ret_type')
319+
src_start = n.properties.get('src_start')
320+
321+
f = Function(path, name, doc, ret_type, src, src_start, src_end)
322+
for arg in args:
323+
name = arg[0]
324+
type_ = arg[1]
325+
f.add_argument(name, type_)
326+
327+
f.id = n.id
328+
329+
return f
330+
306331
def add_function(self, func: Function) -> None:
307332
"""
308333
Adds a function node to the graph database.
@@ -314,7 +339,7 @@ def add_function(self, func: Function) -> None:
314339
q = """MERGE (f:Function:Searchable {path: $path, name: $name,
315340
src_start: $src_start, src_end: $src_end})
316341
SET f.args = $args, f.ret_type = $ret_type, f.src = $src, f.doc = $doc
317-
RETURN ID(f)"""
342+
RETURN f"""
318343

319344
# Prepare arguments in a more straightforward manner
320345
args = [[arg.name, arg.type] for arg in func.args]
@@ -329,34 +354,10 @@ def add_function(self, func: Function) -> None:
329354
'ret_type': func.ret_type
330355
}
331356

332-
res = self._query(q, params)
333-
func.id = res.result_set[0][0]
334-
335-
def _function_from_node(self, n: Node) -> Function:
336-
"""
337-
Create a Function from a graph node
338-
"""
339-
340-
src = n.properties.get('src')
341-
doc = n.properties.get('doc')
342-
path = n.properties.get('path')
343-
name = n.properties.get('name')
344-
args = n.properties.get('args')
345-
src_end = n.properties.get('src_end')
346-
ret_type = n.properties.get('ret_type')
347-
src_start = n.properties.get('src_start')
348-
349-
f = Function(path, name, doc, ret_type, src, src_start, src_end)
350-
for arg in args:
351-
name = arg[0]
352-
type_ = arg[1]
353-
f.add_argument(name, type_)
354-
355-
f.id = n.id
357+
res = self._query(q, params)
358+
node = res.result_set[0][0]
359+
func.id = node.id
356360

357-
return f
358-
359-
360361
# set functions metadata
361362
def set_functions_metadata(self, ids: List[int], metadata: List[dict]) -> None:
362363
assert(len(ids) == len(metadata))
@@ -366,7 +367,8 @@ def set_functions_metadata(self, ids: List[int], metadata: List[dict]) -> None:
366367
WITH $ids[i] AS id, $values[i] AS v
367368
MATCH (f)
368369
WHERE ID(f) = id
369-
SET f += v"""
370+
SET f += v
371+
RETURN f"""
370372

371373
params = {'ids': ids, 'values': metadata}
372374

@@ -483,11 +485,12 @@ def add_file(self, file: File) -> None:
483485
"""
484486

485487
q = """MERGE (f:File:Searchable {path: $path, name: $name, ext: $ext})
486-
RETURN ID(f)"""
488+
RETURN f"""
487489
params = {'path': file.path, 'name': file.name, 'ext': file.ext}
488490

489-
res = self._query(q, params)
490-
file.id = res.result_set[0][0]
491+
res = self._query(q, params)
492+
node = res.result_set[0][0]
493+
file.id = node.id
491494

492495
def delete_files(self, files: List[dict]) -> tuple[str, dict, List[int]]:
493496
"""
@@ -580,7 +583,8 @@ def connect_entities(self, relation: str, src_id: int, dest_id: int) -> None:
580583

581584
q = f"""MATCH (src), (dest)
582585
WHERE ID(src) = $src_id AND ID(dest) = $dest_id
583-
MERGE (src)-[:{relation}]->(dest)"""
586+
MERGE (src)-[e:{relation}]->(dest)
587+
RETURN e"""
584588

585589
params = {'src_id': src_id, 'dest_id': dest_id}
586590
self._query(q, params)
@@ -597,36 +601,12 @@ def function_calls_function(self, caller_id: int, callee_id: int, pos: int) -> N
597601

598602
q = """MATCH (caller:Function), (callee:Function)
599603
WHERE ID(caller) = $caller_id AND ID(callee) = $callee_id
600-
MERGE (caller)-[e:CALLS {pos:$pos}]->(callee)"""
604+
MERGE (caller)-[e:CALLS {pos:$pos}]->(callee)
605+
RETURN e"""
601606

602607
params = {'caller_id': caller_id, 'callee_id': callee_id, 'pos': pos}
603608
self._query(q, params)
604609

605-
def add_struct(self, s: Struct) -> None:
606-
"""
607-
Adds a struct node to the graph database.
608-
609-
Args:
610-
s (Struct): The Struct object to be added.
611-
"""
612-
613-
q = """MERGE (s:Struct:Searchable {name: $name, path: $path, src_start: $src_start,
614-
src_end: $src_end})
615-
SET s.doc = $doc, s.fields = $fields
616-
RETURN ID(s)"""
617-
618-
params = {
619-
'doc': s.doc,
620-
'name': s.name,
621-
'path': s.path,
622-
'src_start': s.src_start,
623-
'src_end': s.src_end,
624-
'fields': s.fields
625-
}
626-
627-
res = self._query(q, params)
628-
s.id = res.result_set[0][0]
629-
630610
def _struct_from_node(self, n: Node) -> Struct:
631611
"""
632612
Create a Struct from a graph node
@@ -652,6 +632,32 @@ def _struct_from_node(self, n: Node) -> Struct:
652632

653633
return s
654634

635+
def add_struct(self, s: Struct) -> None:
636+
"""
637+
Adds a struct node to the graph database.
638+
639+
Args:
640+
s (Struct): The Struct object to be added.
641+
"""
642+
643+
q = """MERGE (s:Struct:Searchable {name: $name, path: $path, src_start: $src_start,
644+
src_end: $src_end})
645+
SET s.doc = $doc, s.fields = $fields
646+
RETURN s"""
647+
648+
params = {
649+
'doc': s.doc,
650+
'name': s.name,
651+
'path': s.path,
652+
'src_start': s.src_start,
653+
'src_end': s.src_end,
654+
'fields': s.fields
655+
}
656+
657+
res = self._query(q, params)
658+
node = res.result_set[0][0]
659+
s.id = node.id
660+
655661
def get_struct_by_name(self, struct_name: str) -> Optional[Struct]:
656662
q = "MATCH (s:Struct) WHERE s.name = $name RETURN s LIMIT 1"
657663
res = self._query(q, {'name': struct_name}).result_set

0 commit comments

Comments
 (0)