tools/update_masked_docs.py

"""This script updates the file torch/_masked/_docs.py that contains
the generated doc-strings for various masked operations. The update
should be triggered whenever a new masked operation is introduced to
torch._masked package. Running the script requires that torch package
is functional.
"""

import os

def main() -> None:

    target = os.path.join('torch', '_masked', '_docs.py')

    try:
        import torch
    except ImportError as msg:
        print(f'Failed to import torch required to build {target}: {msg}')
        return

    if os.path.isfile(target):
        with open(target) as _f:
            current_content = _f.read()
    else:
        current_content = ''

    _new_content = []
    _new_content.append('''\
# -*- coding: utf-8 -*-
# This file is generated, do not modify it!
#
# To update this file, run the update masked docs script as follows:
#
#   python tools/update_masked_docs.py
#
# The script must be called from an environment where the development
# version of torch package can be imported and is functional.
#
''')

    for func_name in sorted(torch._masked.__all__):
        func = getattr(torch._masked, func_name)
        func_doc = torch._masked._generate_docstring(func)
        _new_content.append(f'{func_name}_docstring = """{func_doc}"""\n')

    new_content = '\n'.join(_new_content)

    if new_content == current_content:
        print(f'Nothing to update in {target}')
        return

    with open(target, 'w') as _f:
        _f.write(new_content)

    print(f'Successfully updated {target}')


if __name__ == '__main__':
    main()