diff --git a/src/codesnap/core/complexity.py b/src/codesnap/core/complexity.py new file mode 100644 index 0000000..a823cca --- /dev/null +++ b/src/codesnap/core/complexity.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +import ast +from codesnap.core.models import FunctionInfo + + +class ComplexityCalculator: + def calculate_function_complexity(self, node: ast.FunctionDef | ast.AsyncFunctionDef, source: str) -> FunctionInfo: + lines = source.split("\n") + start_line = node.lineno + end_line = node.end_lineno or node.lineno + + body = lines[start_line - 1 : end_line] + code = "\n".join(body) + + complexity = 1 + for child in ast.walk(node): + if isinstance(child, (ast.If, ast.For, ast.While, ast.Try, ast.With, ast.AsyncWith)): + complexity += 1 + elif isinstance(child, ast.BoolOp): + complexity += len(child.values) - 1 + elif isinstance(child, ast.comprehension): + complexity += 1 + + params = len(node.args.args) + + docstring = ast.get_docstring(node) + + return FunctionInfo( + name=node.name, + complexity=complexity, + parameters=params, + lines=end_line - start_line + 1, + code=code, + docstring=docstring, + )