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)