diff --git a/litedoc/syntax/node.py b/litedoc/syntax/node.py index d56b039..04b101a 100644 --- a/litedoc/syntax/node.py +++ b/litedoc/syntax/node.py @@ -14,6 +14,7 @@ from litedoc.docstring.docstring import Docstring from litedoc.i18n import get_text, litedoc_hide +from litedoc.utils import remove_docstrings_from_code class TypeHint: @@ -146,7 +147,7 @@ class FunctionNode(BaseModel): return_: str = TypeHint.NO_RETURN decorators: list[str] = [] - src: str + src: str # 源码 is_async: bool = False is_classmethod: bool = False @@ -294,7 +295,8 @@ def markdown(self, lang: str, indent: int = 0, **kwargs) -> str: md += f"{self.docs.markdown(lang, indent)}\n" else: pass - # 源码展示 + + """源码github链接""" if kwargs.get("bu", None): # 源码链接 self.module_file_path = self.module_file_path.replace("\\", "/") @@ -304,7 +306,9 @@ def markdown(self, lang: str, indent: int = 0, **kwargs) -> str: or_and_a = f" {get_text(lang, 'or')} {a_tag}" else: or_and_a = "" - md += PREFIX + f"\n
\n {get_text(lang, 'src')}{or_and_a}\n\n```python\n{self.src}\n```\n
\n\n" + + """源码展示""" + md += PREFIX + f"\n
\n {get_text(lang, 'src')}{or_and_a}\n\n```python\n{self.get_src_without_docstring()}\n```\n
\n\n" return md @@ -317,6 +321,14 @@ def complete_default_args(self): num = len(self.args) + len(self.posonlyargs) - len(self.defaults) self.defaults = [ConstantNode(value=TypeHint.NO_DEFAULT) for _ in range(num)] + self.defaults + def get_src_without_docstring(self): + """ + 获取去掉docstring的源码 + Returns: + str + """ + return remove_docstrings_from_code(self.src) + def __str__(self): return f"def {self.name}({', '.join([f'{arg.name}: {arg.type} = {arg.default}' for arg in self.args])}) -> {self.return_}" diff --git a/litedoc/utils.py b/litedoc/utils.py new file mode 100644 index 0000000..675e13d --- /dev/null +++ b/litedoc/utils.py @@ -0,0 +1,15 @@ +import ast + + +def remove_docstring_from_function(node): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + if len(node.body) > 0 and isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value, ast.Str): + node.body.pop(0) + return node + + +def remove_docstrings_from_code(source_code): + tree = ast.parse(source_code) + for node in ast.walk(tree): + remove_docstring_from_function(node) + return ast.unparse(tree)