Add utilities in pg.object_utils for handling docstrs.

PiperOrigin-RevId: 523145800

.github/workflows/ci.yaml

@@ -27,6 +27,7 @@ jobs:
         pip install pytest
         pip install pytest-xdist
         pip install pytest-cov
+        pip install -r requirements.txt
     - name: Test with pytest and generate coverage report
       run: |
         pytest -n auto --cov=pyglove --cov-report=xml

.github/workflows/docgen_test.yaml

@@ -25,6 +25,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
+        pip install -r requirements.txt
         pip install -r docs/requirements.txt
     - name: Print installed dependencies
       run: |

.github/workflows/pypi-nightly.yaml

@@ -22,6 +22,7 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         pip install setuptools wheel twine
+        pip install -r requirements.txt
     - name: Build and publish
       env:
         TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}

.github/workflows/pypi.yaml

@@ -20,6 +20,7 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         pip install setuptools wheel twine
+        pip install -r requirements.txt
     - name: Build and publish
       env:
         TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}

README.md

@@ -22,7 +22,7 @@ direct manipulation of objects that makes meta-programs much easier to write.
 It has been used to handle complex machine learning scenarios, such as AutoML,
 as well as facilitating daily programming tasks with extra flexibility.
 
-PyGlove is lightweight and has no dependencies beyond the Python interpreter.
+PyGlove is lightweight and has very few dependencies beyond the Python interpreter.
 It provides:
 
 * A mutable symbolic object model for Python;

pyglove/core/__init__.py

@@ -263,12 +263,13 @@
 Formattable = object_utils.Formattable
 MaybePartial = object_utils.MaybePartial
 JSONConvertible = object_utils.JSONConvertible
+DocStr = object_utils.DocStr
 
 registered_types = object_utils.registered_types
 is_partial = object_utils.is_partial
 format = object_utils.format   # pylint: disable=redefined-builtin
 print = object_utils.print    # pylint: disable=redefined-builtin
-
+docstr = object_utils.docstr
 
 #
 # Symbols from `logging.py`.

pyglove/core/object_utils/__init__.py

@@ -12,13 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # pylint: disable=line-too-long
-"""Utility library for handling hierarchical Python objects.
+"""Utility library that provides common traits for objects in Python.
 
 Overview
 --------
 
-``pg.object_utils`` facilitates the handling of hierarchical
-Python objects. It sits at the bottom of all PyGlove modules and empowers other
+``pg.object_utils`` sits at the bottom of all PyGlove modules and empowers other
 modules with the following features:
 
   +---------------------+--------------------------------------------+
@@ -62,6 +61,10 @@
   |                     |                                            |
   |                     | :func:`pg.object_utils.flatten`            |
   +---------------------+--------------------------------------------+
+  | Code generation     | :class:`pg.object_utils.make_function`     |
+  +---------------------+--------------------------------------------+
+  | Docstr handling     | :class:`pg.docstr`,                        |
+  +---------------------+--------------------------------------------+
 """
 # pylint: enable=line-too-long
 # pylint: disable=g-bad-import-order
@@ -112,5 +115,14 @@
 # Handling code generation.
 from pyglove.core.object_utils.codegen import make_function
 
+# Handling docstrings.
+from pyglove.core.object_utils.docstr_utils import DocStr
+from pyglove.core.object_utils.docstr_utils import DocStrStyle
+from pyglove.core.object_utils.docstr_utils import DocStrEntry
+from pyglove.core.object_utils.docstr_utils import DocStrExample
+from pyglove.core.object_utils.docstr_utils import DocStrArgument
+from pyglove.core.object_utils.docstr_utils import DocStrReturns
+from pyglove.core.object_utils.docstr_utils import DocStrRaises
+from pyglove.core.object_utils.docstr_utils import docstr
 
 # pylint: enable=g-bad-import-order

pyglove/core/object_utils/docstr_utils.py

@@ -0,0 +1,139 @@
+# Copyright 2023 The PyGlove Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Utilities for working with docstrs."""
+
+import dataclasses
+import enum
+from typing import Any, Dict, List, Optional
+
+import docstring_parser
+
+
+class DocStrStyle(enum.Enum):
+  """Docstring style."""
+  REST = 1
+  GOOGLE = 2
+  NUMPYDOC = 3
+  EPYDOC = 4
+
+
+@dataclasses.dataclass
+class DocStrEntry:
+  """An entry in a docstring."""
+  description: str
+
+
+@dataclasses.dataclass
+class DocStrArgument(DocStrEntry):
+  """An entry in the "Args" section of a docstring."""
+  name: str
+  type_name: Optional[str] = None
+  default: Optional[str] = None
+  is_optional: Optional[bool] = None
+
+
+@dataclasses.dataclass
+class DocStrReturns(DocStrEntry):
+  """An entry in the "Returns"/"Yields" section of a docstring."""
+  name: Optional[str] = None
+  type_name: Optional[str] = None
+  is_yield: bool = False
+
+
+@dataclasses.dataclass
+class DocStrRaises(DocStrEntry):
+  """An entry in the "Raises" section of a docstring."""
+  type_name: Optional[str] = None
+
+
+@dataclasses.dataclass
+class DocStrExample(DocStrEntry):
+  """An entry in the "Examples" section of a docstring."""
+
+
+@dataclasses.dataclass
+class DocStr:
+  """Docstring."""
+  style: DocStrStyle
+  short_description: Optional[str]
+  long_description: Optional[str]
+  examples: List[DocStrExample]
+  args: Dict[str, DocStrArgument]
+  returns: Optional[DocStrReturns]
+  raises: List[DocStrRaises]
+  blank_after_short_description: bool = True
+
+  @classmethod
+  def parse(cls, text: str, style: Optional[DocStrStyle] = None) -> 'DocStr':
+    """Parses a docstring."""
+    result = docstring_parser.parse(text, _to_parser_style(style))
+    return cls(
+        style=_from_parser_style(result.style),
+        short_description=result.short_description,
+        long_description=result.long_description,
+        examples=[
+            DocStrExample(description=e.description)
+            for e in result.examples
+        ],
+        args={  # pylint: disable=g-complex-comprehension
+            p.arg_name: DocStrArgument(
+                name=p.arg_name, description=p.description,
+                type_name=p.type_name, default=p.default,
+                is_optional=p.is_optional)
+            for p in result.params
+        },
+        returns=DocStrReturns(  # pylint: disable=g-long-ternary
+            name=result.returns.return_name,
+            description=result.returns.description,
+            is_yield=result.returns.is_generator) if result.returns else None,
+        raises=[
+            DocStrRaises(type_name=r.type_name, description=r.description)
+            for r in result.raises
+        ],
+        blank_after_short_description=result.blank_after_short_description)
+
+
+def docstr(symbol: Any) -> Optional[DocStr]:
+  """Gets structure docstring of a Python symbol."""
+  docstr_text = getattr(symbol, '__doc__', None)
+  return DocStr.parse(docstr_text) if docstr_text else None
+
+
+_PARSER_STYLE_MAPPING = [
+    (DocStrStyle.REST, docstring_parser.DocstringStyle.REST),
+    (DocStrStyle.GOOGLE, docstring_parser.DocstringStyle.GOOGLE),
+    (DocStrStyle.NUMPYDOC, docstring_parser.DocstringStyle.NUMPYDOC),
+    (DocStrStyle.EPYDOC, docstring_parser.DocstringStyle.EPYDOC),
+]
+
+
+def _to_parser_style(
+    style: Optional[DocStrStyle]) -> docstring_parser.DocstringStyle:
+  """Returns parser style from DocStrStyle."""
+  if style is None:
+    return docstring_parser.DocstringStyle.AUTO
+  for s, ps in _PARSER_STYLE_MAPPING:
+    if style == s:
+      return ps
+  raise ValueError(f'Unsupported style {style}.')
+
+
+def _from_parser_style(
+    style: docstring_parser.DocstringStyle) -> DocStrStyle:
+  """Returns DocStrStyle from parser style."""
+  for s, ps in _PARSER_STYLE_MAPPING:
+    if style == ps:
+      return s
+  raise ValueError(f'Unsupported parser style {style}.')
+

pyglove/core/object_utils/docstr_utils_test.py

@@ -0,0 +1,105 @@
+# Copyright 2023 The PyGlove Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for pyglove.object_utils.docstr_utils."""
+
+import unittest
+from pyglove.core.object_utils import docstr_utils
+
+
+class DocStrTest(unittest.TestCase):
+  """Tests for DocStr."""
+
+  def test_parse(self):
+    doc = """Test doc string."""
+    docstr = docstr_utils.DocStr.parse(doc)
+    # Returns the first style that fit.
+    self.assertEqual(docstr.style, docstr_utils.DocStrStyle.REST)
+
+    docstr = docstr_utils.DocStr.parse(doc, docstr_utils.DocStrStyle.GOOGLE)
+    self.assertEqual(docstr.style, docstr_utils.DocStrStyle.GOOGLE)
+
+  def test_docstr(self):
+    def my_sum(x, y, *args, **kwargs):
+      """Returns the sum of two integers.
+
+      This function will return the sum of two integers.
+
+      Examples:
+
+        ```
+          ret = sum(1, 2)
+          print(ret)
+        ```
+
+      Args:
+        x: An integer.
+        y: Another integer.
+        *args: Variable positional args.
+        **kwargs: Variable keyword args.
+
+      Returns:
+        The sum of both.
+
+      Raises:
+        ValueError: when either `x` and `y` is not an integer.
+      """
+      del args, kwargs
+      return x + y
+
+    self.assertEqual(docstr_utils.docstr(my_sum), docstr_utils.DocStr(
+        style=docstr_utils.DocStrStyle.GOOGLE,
+        short_description='Returns the sum of two integers.',
+        long_description='This function will return the sum of two integers.',
+        examples=[
+            docstr_utils.DocStrExample(
+                description='```\n  ret = sum(1, 2)\n  print(ret)\n```')
+        ],
+        args={
+            'x': docstr_utils.DocStrArgument(
+                name='x',
+                description='An integer.',
+            ),
+            'y': docstr_utils.DocStrArgument(
+                name='y',
+                description='Another integer.',
+            ),
+            '*args': docstr_utils.DocStrArgument(
+                name='*args',
+                description='Variable positional args.',
+            ),
+            '**kwargs': docstr_utils.DocStrArgument(
+                name='**kwargs',
+                description='Variable keyword args.',
+            )
+        },
+        returns=docstr_utils.DocStrReturns(
+            description='The sum of both.',
+        ),
+        raises=[
+            docstr_utils.DocStrRaises(
+                type_name='ValueError',
+                description='when either `x` and `y` is not an integer.',
+            )
+        ]
+    ))
+
+    class Foo:
+      pass
+
+    self.assertIsNone(docstr_utils.docstr(Foo))
+    self.assertIsNone(docstr_utils.docstr(None))
+
+
+if __name__ == '__main__':
+  unittest.main()

requirements.txt

@@ -0,0 +1 @@
+docstring-parser>=0.12

setup.py

@@ -38,9 +38,22 @@ def _get_version():
   return version
 
 
-def _parse_requirements(requirements_txt_path):
+def _parse_requirements(requirements_txt_path: str) -> list[str]:
+  """Returns a list of dependencies for setup() from requirements.txt."""
+
+  def _strip_comments_from_line(s: str) -> str:
+    """Parses a line of a requirements.txt file."""
+    requirement, *_ = s.split('#')
+    return requirement.strip()
+
+  # Currently a requirements.txt is being used to specify dependencies. In order
+  # to avoid specifying it in two places, we're going to use that file as the
+  # source of truth.
   with open(requirements_txt_path) as fp:
-    return fp.read().splitlines()
+    # Parse comments.
+    lines = [_strip_comments_from_line(line) for line in fp.read().splitlines()]
+    # Remove empty lines and direct github repos (not allowed in PyPI setups)
+    return [l for l in lines if (l and 'github.com' not in l)]
 
 
 _VERSION = _get_version()
@@ -57,7 +70,7 @@ def _parse_requirements(requirements_txt_path):
     author_email='[email protected]',
     # Contained modules and scripts.
     packages=find_namespace_packages(include=['pyglove*']),
-    install_requires=[],
+    install_requires=_parse_requirements('requirements.txt'),
     extras_require={},
     requires_python='>=3.7',
     include_package_data=True,
@@ -76,5 +89,7 @@ def _parse_requirements(requirements_txt_path):
         'Topic :: Software Development :: Libraries :: Python Modules',
         'Topic :: Software Development :: Libraries',
     ],
-    keywords='ai machine learning automl mutable symbolic framework meta-programming',
+    keywords=(
+        'ai machine learning automl mutable symbolic '
+        'framework meta-programming'),
 )