Coverage for lobster/tools/python/python.py: 0%
338 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-08-05 10:00 +0000
« prev ^ index » next coverage.py v7.10.7, created at 2026-08-05 10:00 +0000
1#!/usr/bin/env python3
2#
3# lobster_python - Extract Python tracing tags for LOBSTER
4# Copyright (C) 2022-2023 Bayerische Motoren Werke Aktiengesellschaft (BMW AG)
5#
6# This program is free software: you can redistribute it and/or modify
7# it under the terms of the GNU Affero General Public License as
8# published by the Free Software Foundation, either version 3 of the
9# License, or (at your option) any later version.
10#
11# This program is distributed in the hope that it will be useful, but
12# WITHOUT ANY WARRANTY; without even the implied warranty of
13# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
14# Affero General Public License for more details.
15#
16# You should have received a copy of the GNU Affero General Public
17# License along with this program. If not, see
18# <https://www.gnu.org/licenses/>.
20from argparse import Namespace
21import sys
22import os.path
23import multiprocessing
24import functools
25import re
26from dataclasses import dataclass
27from typing import List, Optional, Sequence, Tuple
29from libcst.metadata import PositionProvider
30import libcst as cst
32from lobster.common.items import Tracing_Tag, Implementation, Activity
33from lobster.common.location import File_Reference
34from lobster.common.io import lobster_write, ensure_output_directory
35from lobster.common.meta_data_tool_base import MetaDataToolBase
37LOBSTER_TRACE_PREFIX = "# lobster-trace: "
38LOBSTER_JUST_PREFIX = "# lobster-exclude: "
39func_name = []
42@dataclass
43class PythonToolConfig:
44 files: List[str]
45 activity: bool = False
46 out: Optional[str] = None
47 single: bool = False
48 only_tagged_functions: bool = False
49 parse_decorator: Optional[Tuple[str, str]] = None
50 parse_versioned_decorator: Optional[Tuple[str, str, str]] = None
53def count_occurrence_of_last_function_from_function_name_list(function_names):
54 """
55 Returns the last function and class name (if present) in a list along with
56 the count of its previous occurrences.
58 The function identifies the last entry in the `function_names` list, extracts
59 the function and class names (if applicable), and counts prior occurrences of
60 the same function.
61 The result is formatted as `module.class.function-count` or `module.function-count`.
63 Args:
64 function_names (list):
65 List of strings formatted as `module.class.function:line_number`
66 or `module.function:line_number`.
68 Returns:
69 str: The last function (and class if applicable) with its occurrence count,
70 formatted as `module.class.function-count` or `module.function-count`.
72 Examples:
73 function_names = ['hello.add:2', 'hello.sub:5', 'hello.add:8']
74 returns: 'hello.add-2'
75 class_function_names = ['Example.hello.add:2', 'Example.hello.sub:5',]
76 returns: 'Example.hello.add-2'
77 """
78 function_and_file_name = re.split(r"[.:]", function_names[-1])
79 class_name_with_module = function_names[-1].split(':', 1)[0].split(".")
81 if len(class_name_with_module) == 3:
82 function_and_file_name[1] = (class_name_with_module[1] + '.' +
83 class_name_with_module[2])
85 filename = function_and_file_name[0]
86 last_function = function_and_file_name[1]
87 count = 0
88 for element in range(0, len(function_names) - 1):
89 class_name_with_function = function_names[element].split(':', 1)[0].split(".")
90 if len(class_name_with_function) == 3:
91 if last_function == (class_name_with_function[1] + '.' +
92 class_name_with_function[2]):
93 count += 1
94 if re.split(r"[.:]", function_names[element])[1] == last_function:
95 count += 1
96 function_name = (filename + "." + last_function +
97 ("-" + str(count) if count > 0 else ''))
99 return function_name
102def parse_value(val):
103 if isinstance(val, cst.SimpleString):
104 return val.value[1:-1]
105 if isinstance(val, cst.List):
106 return [parse_value(item.value)
107 for item in val.elements]
109 rv = str(val.value)
110 if rv == "None":
111 rv = None
112 return rv
115class Python_Traceable_Node:
116 def __init__(self, location, name, kind):
117 assert isinstance(location, File_Reference)
118 assert isinstance(name, str)
119 assert isinstance(kind, str)
120 self.location = location
121 self.name = name
122 self.kind = kind
123 self.parent = None
124 self.children = []
125 self.tags = []
126 self.just = []
128 def register_tag(self, tag):
129 assert isinstance(tag, Tracing_Tag)
130 self.tags.append(tag)
132 def register_justification(self, justification):
133 assert isinstance(justification, str)
134 self.just.append(justification)
136 def set_parent(self, node):
137 assert isinstance(node, Python_Traceable_Node)
138 node.children.append(self)
139 self.parent = node
141 def to_json(self):
142 return {"kind" : self.kind,
143 "name" : self.name,
144 "tags" : [x.to_json() for x in self.tags],
145 "just" : self.just,
146 "children" : [x.to_json() for x in self.children]}
148 def to_lobster(self, schema, items):
149 assert schema is Implementation or schema is Activity
150 assert isinstance(items, list)
151 assert False
153 def fqn(self):
154 if self.parent:
155 rv = self.parent.fqn() + "."
156 else:
157 rv = ""
158 if self.location.line is not None and \
159 isinstance(self, Python_Function):
160 rv += f"{self.name}:{str(self.location.line)}"
161 else:
162 rv += self.name
163 return rv
165 def lobster_tag(self):
166 return Tracing_Tag("python", self.fqn())
168 def warn_ignored(self, reason):
169 for tag in self.tags:
170 print(f"{self.location.to_string()}: warning: ignored tag {tag}"
171 f" because {reason} already has annotations")
172 for just in self.just:
173 print(f"{self.location.to_string()}: warning: "
174 f"ignored justification '{just}' "
175 f"because {reason} already has annotations")
178class Python_Module(Python_Traceable_Node):
179 def __init__(self, location, name):
180 super().__init__(location, name, "Module")
182 def to_lobster(self, schema, items):
183 assert schema is Implementation or schema is Activity
184 assert isinstance(items, list)
185 for node in self.children:
186 node.to_lobster(schema, items)
189class Python_Class(Python_Traceable_Node):
190 def __init__(self, location, name):
191 super().__init__(location, name, "Class")
193 def to_lobster(self, schema, items):
194 assert schema is Implementation or schema is Activity
195 assert isinstance(items, list)
196 # Classes are dealt with a bit differently. If you add a tag
197 # or justification to a class, then children are ignored, and
198 # we trace to the class.
199 #
200 # Alternatively, can leave out the tag and instead trace to
201 # each child.
203 # First get child items
204 class_contents = []
205 for node in self.children:
206 node.to_lobster(schema, class_contents)
208 # If we're extracting pyunit/unittest items, then we always ignore
209 # classes, but we do add our tags to all the tests.
210 if schema is Activity:
211 for item in class_contents:
212 for tag in self.tags:
213 item.add_tracing_target(tag)
214 items += class_contents
215 return
217 l_item = Implementation(tag = Tracing_Tag("python",
218 self.fqn()),
219 location = self.location,
220 language = "Python",
221 kind = self.kind,
222 name = self.fqn())
224 # If we have tags or justifications on the class itself, we
225 # give precedence to that.
226 if self.tags or self.just:
227 for tag in self.tags:
228 l_item.add_tracing_target(tag)
229 l_item.just_up += self.just
231 for c_item in self.children:
232 c_item.warn_ignored(self.name)
234 items.append(l_item)
235 return
237 # Otherwise, we ignore the class and instead trace to each
238 # child
239 items += class_contents
242class Python_Function(Python_Traceable_Node):
243 def __init__(self, location, name):
244 super().__init__(location, name, "Function")
246 def set_parent(self, node):
247 assert isinstance(node, Python_Traceable_Node)
248 node.children.append(self)
249 self.parent = node
250 if isinstance(node, Python_Class):
251 if self.name == "__init__":
252 self.kind = "Constructor"
253 else:
254 self.kind = "Method"
256 def to_lobster(self, schema, items):
257 assert schema is Implementation or schema is Activity
258 assert isinstance(items, list)
260 func_name.append(self.fqn())
261 tagname = count_occurrence_of_last_function_from_function_name_list(
262 func_name
263 )
264 pattern = r"[-]"
265 val = re.split(pattern, tagname)
266 name_value = val[0]
268 if schema is Implementation:
269 l_item = Implementation(tag = Tracing_Tag("python",
270 tagname),
271 location = self.location,
272 language = "Python",
273 kind = self.kind,
274 name = name_value)
275 elif self.name.startswith("test") or self.name.startswith("_test") \
276 or self.name.endswith("test"):
277 l_item = Activity(tag = Tracing_Tag("pyunit",
278 self.fqn()),
279 location = self.location,
280 framework = "PyUnit",
281 kind = "Test")
282 else:
283 return
285 for tag in self.tags:
286 l_item.add_tracing_target(tag)
287 l_item.just_up += self.just
289 # Any children of functions are not testable units. Their
290 # tracing tags contribute to ours, but otherwise they don't
291 # appear.
292 nested_items = []
293 for node in self.children:
294 node.to_lobster(schema, nested_items)
295 for item in nested_items:
296 # TODO: Warn about useless nested justifications
297 # Merge tracing tags
298 for tag in item.unresolved_references:
299 l_item.add_tracing_target(tag)
301 items.append(l_item)
304class Lobster_Visitor(cst.CSTVisitor):
305 METADATA_DEPENDENCIES = (PositionProvider,)
307 def __init__(self, file_name, options):
308 super().__init__()
309 assert os.path.isfile(file_name)
310 self.file_name = file_name
312 self.module = Python_Module(
313 File_Reference(file_name),
314 os.path.basename(file_name).replace(".py", ""))
316 self.activity = options["activity"]
317 self.current_node = None
318 self.stack = [self.module]
320 self.namespace = options["namespace"]
321 self.exclude_untagged = options["exclude_untagged"]
323 self.decorator_name = options["decorator"]
324 self.dec_arg_name = options["dec_arg_name"]
325 self.dec_arg_version = options["dec_arg_version"]
327 def parse_dotted_name(self, name):
328 if isinstance(name, cst.Call):
329 return self.parse_dotted_name(name.func)
330 if isinstance(name, cst.Name):
331 return name.value
332 if isinstance(name, cst.Attribute):
333 # value -- prefix
334 # attr -- postfix
335 return f"{self.parse_dotted_name(name.value)}." \
336 f"{self.parse_dotted_name(name.attr)}"
337 return None
339 def parse_decorators(self, decorators):
340 for dec in decorators:
341 dec_name = self.parse_dotted_name(dec.decorator)
342 if dec_name is None:
343 continue
344 if dec_name != self.decorator_name:
345 continue
346 dec_args = {arg.keyword.value: parse_value(arg.value)
347 for arg in dec.decorator.args}
349 # TODO: Better error messages if these assumptions are
350 # violated
351 assert self.dec_arg_name in dec_args
352 if self.dec_arg_version:
353 assert self.dec_arg_version in dec_args
354 tag = Tracing_Tag(self.namespace,
355 dec_args[self.dec_arg_name],
356 dec_args.get(self.dec_arg_version, None))
357 self.current_node.register_tag(tag)
359 elif isinstance(dec_args[self.dec_arg_name], list):
360 for item in dec_args[self.dec_arg_name]:
361 tag = Tracing_Tag(self.namespace, item)
362 self.current_node.register_tag(tag)
364 else:
365 tag = Tracing_Tag(self.namespace,
366 dec_args[self.dec_arg_name])
367 self.current_node.register_tag(tag)
369 def visit_ClassDef(self, node):
370 line = self.get_metadata(PositionProvider, node).start.line
371 loc = File_Reference(self.file_name, line)
372 t_item = Python_Class(loc, node.name.value)
373 t_item.set_parent(self.stack[-1])
374 self.stack.append(t_item)
375 self.current_node = t_item
376 self.parse_decorators(node.decorators)
378 def visit_FunctionDef(self, node):
379 line = self.get_metadata(PositionProvider, node).start.line
380 loc = File_Reference(self.file_name, line)
381 t_item = Python_Function(loc, node.name.value)
382 t_item.set_parent(self.stack[-1])
383 self.stack.append(t_item)
384 self.current_node = t_item
385 self.parse_decorators(node.decorators)
387 def leave_FunctionDef(self, original_node):
388 self.stack.pop()
389 self.current_node = self.stack[-1]
391 def leave_ClassDef(self, original_node):
392 self.stack.pop()
393 self.current_node = self.stack[-1]
395 def visit_Comment(self, node):
396 line = self.get_metadata(PositionProvider, node).start.line
397 # For some reason the comment in a class is associated with
398 # its constructor. We can check if it preceeds it (by line),
399 # and so associate it with the enclosing item.
400 if self.current_node and \
401 self.current_node.location.line and \
402 self.current_node.location.line > line:
403 actual = self.current_node.parent
404 else:
405 actual = self.current_node
407 if node.value.startswith(LOBSTER_TRACE_PREFIX):
408 tag = node.value[len(LOBSTER_TRACE_PREFIX):].strip()
409 actual.register_tag(
410 Tracing_Tag.from_text(self.namespace,
411 tag))
413 elif node.value.startswith(LOBSTER_JUST_PREFIX):
414 reason = node.value[len(LOBSTER_JUST_PREFIX):].strip()
415 actual.register_justification(reason)
418def process_file(file_name, options):
419 # pylint: disable=protected-access
420 assert isinstance(file_name, str)
421 assert isinstance(options, dict)
423 items = []
424 try:
425 with open(file_name, encoding="UTF-8") as fd:
426 ast = cst.parse_module(fd.read())
428 ast = cst.MetadataWrapper(ast)
429 visitor = Lobster_Visitor(file_name, options)
430 ast.visit(visitor)
432 if options["activity"]:
433 visitor.module.to_lobster(Activity, items)
434 else:
435 visitor.module.to_lobster(Implementation, items)
437 if options["exclude_untagged"]:
438 items = [item for item in items if item.unresolved_references]
440 return True, items
442 except cst._exceptions.ParserSyntaxError as exc:
443 print(file_name, exc.message)
444 return False, []
446 except UnicodeDecodeError as exc:
447 print(file_name, str(exc))
448 return False, []
450 except Exception as exc:
451 print(f"Unspecified issue in file: {file_name}")
452 raise
455class PythonTool(MetaDataToolBase):
456 def __init__(self):
457 super().__init__(
458 name="python",
459 description="Extract tracing tags from Python code or tests",
460 official=True,
461 )
462 ap = self._argument_parser
463 ap.add_argument("files",
464 nargs="+",
465 metavar="FILE|DIR")
466 ap.add_argument("--activity",
467 action="store_true",
468 default=False,
469 help=("generate activity traces (tests) instead of"
470 " an implementation trace"))
471 ap.add_argument("--out",
472 default=None)
473 ap.add_argument("--single",
474 action="store_true",
475 default=False,
476 help="don't multi-thread")
477 ap.add_argument("--only-tagged-functions",
478 default=False,
479 action="store_true",
480 help="only trace functions with tags")
481 grp = ap.add_mutually_exclusive_group()
482 grp.add_argument("--parse-decorator",
483 nargs=2,
484 metavar=("DECORATOR", "NAME_ARG"),
485 default=(None, None))
486 grp.add_argument("--parse-versioned-decorator",
487 nargs=3,
488 metavar=("DECORATOR", "NAME_ARG", "VERSION_ARG"),
489 default=(None, None, None))
491 def _run_impl(self, options: Namespace) -> int:
492 parse_decorator = (
493 options.parse_decorator
494 if options.parse_decorator[0] is not None
495 else None
496 )
497 parse_versioned_decorator = (
498 options.parse_versioned_decorator
499 if options.parse_versioned_decorator[0] is not None
500 else None
501 )
502 config = PythonToolConfig(
503 files=options.files,
504 activity=options.activity,
505 out=options.out,
506 single=options.single,
507 only_tagged_functions=options.only_tagged_functions,
508 parse_decorator=parse_decorator,
509 parse_versioned_decorator=parse_versioned_decorator,
510 )
511 try:
512 if run_lobster_python(config):
513 return 0
514 except ValueError as exc:
515 self._argument_parser.error(str(exc))
516 return 1
517 print("Note: Earlier parse errors make actual output unreliable")
518 return 1
521def collect_python_files(files: List[str]) -> List[str]:
522 file_list = []
523 for item in files:
524 if os.path.isfile(item):
525 file_list.append(item)
526 elif os.path.isdir(item):
527 for path, _, candidates in os.walk(item):
528 for filename in candidates:
529 _, ext = os.path.splitext(filename)
530 if ext == ".py":
531 file_list.append(os.path.join(path, filename))
532 else:
533 raise ValueError(f"{item} is not a file or directory")
534 return file_list
537def build_context(config: PythonToolConfig) -> dict:
538 context = {
539 "activity" : config.activity,
540 "decorator" : None,
541 "dec_arg_name" : None,
542 "dec_arg_version" : None,
543 "exclude_untagged" : config.only_tagged_functions,
544 "namespace" : "req",
545 }
547 if config.parse_decorator and config.parse_versioned_decorator:
548 raise ValueError(
549 "Only one of parse_decorator or parse_versioned_decorator can be set"
550 )
552 if config.parse_decorator:
553 context["decorator"] = config.parse_decorator[0]
554 context["dec_arg_name"] = config.parse_decorator[1]
555 elif config.parse_versioned_decorator:
556 context["decorator"] = config.parse_versioned_decorator[0]
557 context["dec_arg_name"] = config.parse_versioned_decorator[1]
558 context["dec_arg_version"] = config.parse_versioned_decorator[2]
560 return context
563def run_lobster_python(config: PythonToolConfig) -> bool:
564 file_list = collect_python_files(config.files)
565 context = build_context(config)
567 pfun = functools.partial(process_file, options=context)
568 items = []
569 ok = True
571 if config.single:
572 for file_name in file_list:
573 new_ok, new_items = pfun(file_name)
574 ok &= new_ok
575 items += new_items
576 else:
577 with multiprocessing.Pool() as pool:
578 for new_ok, new_items in pool.imap_unordered(pfun, file_list):
579 ok &= new_ok
580 items += new_items
582 schema = Activity if config.activity else Implementation
584 if config.out:
585 ensure_output_directory(config.out)
586 with open(config.out, "w", encoding="UTF-8") as fd:
587 lobster_write(fd, schema, "lobster_python", items)
588 print(f"Written output for {len(items)} items to {config.out}")
589 else:
590 lobster_write(sys.stdout, schema, "lobster_python", items)
591 print()
593 return ok
596def lobster_python(config: PythonToolConfig) -> None:
597 """This is an API function."""
598 run_lobster_python(config)
601def main(args: Optional[Sequence[str]] = None) -> int:
602 return PythonTool().run(args)