Skip to content

Commit

Permalink
feat: add convert_document_to_annotate_file_json
Browse files Browse the repository at this point in the history
  • Loading branch information
galz10 committed Jun 7, 2023
1 parent 70106c1 commit 221a369
Show file tree
Hide file tree
Showing 3 changed files with 34 additions and 8 deletions.
13 changes: 13 additions & 0 deletions google/cloud/documentai_toolbox/wrappers/document.py
Original file line number Diff line number Diff line change
Expand Up @@ -732,6 +732,19 @@ def convert_document_to_annotate_file_response(self) -> AnnotateFileResponse:
"""
return _convert_to_vision_annotate_file_response(self.text, self.pages)

def convert_document_to_annotate_file_json_response(self) -> str:
r"""Convert OCR data from Document.proto to JSON str of AnnotateFileResponse for Vision API.
Args:
None.
Returns:
str:
JSON string of TextAnnotations.
"""
return AnnotateFileResponse.to_json(
_convert_to_vision_annotate_file_response(self.text, self.pages)
)

def export_images(
self, output_path: str, output_file_prefix: str, output_file_extension: str
) -> List[str]:
Expand Down
16 changes: 8 additions & 8 deletions scripts/readme-gen/readme_gen.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,37 +33,37 @@
autoescape=True,
)

README_TMPL = jinja_env.get_template('README.tmpl.rst')
README_TMPL = jinja_env.get_template("README.tmpl.rst")


def get_help(file):
return subprocess.check_output(['python', file, '--help']).decode()
return subprocess.check_output(["python", file, "--help"]).decode()


def main():
parser = argparse.ArgumentParser()
parser.add_argument('source')
parser.add_argument('--destination', default='README.rst')
parser.add_argument("source")
parser.add_argument("--destination", default="README.rst")

args = parser.parse_args()

source = os.path.abspath(args.source)
root = os.path.dirname(source)
destination = os.path.join(root, args.destination)

jinja_env.globals['get_help'] = get_help
jinja_env.globals["get_help"] = get_help

with io.open(source, 'r') as f:
with io.open(source, "r") as f:
config = yaml.load(f)

# This allows get_help to execute in the right directory.
os.chdir(root)

output = README_TMPL.render(config)

with io.open(destination, 'w') as f:
with io.open(destination, "w") as f:
f.write(output)


if __name__ == '__main__':
if __name__ == "__main__":
main()
13 changes: 13 additions & 0 deletions tests/unit/test_document.py
Original file line number Diff line number Diff line change
Expand Up @@ -536,6 +536,19 @@ def test_convert_document_to_annotate_file_response():
assert actual == expected


def test_convert_document_to_annotate_file_json_response():
doc = document.Document.from_document_path(
document_path="tests/unit/resources/0/toolbox_invoice_test-0.json"
)

actual = doc.convert_document_to_annotate_file_json_response()

with open("tests/unit/resources/toolbox_invoice_test-0-vision.json", "r") as f:
expected = f.read()

assert actual == expected


def test_export_images(get_bytes_images_mock):
doc = document.Document.from_gcs(
gcs_bucket_name="test-directory", gcs_prefix="documentai/output/123456789/0"
Expand Down

0 comments on commit 221a369

Please sign in to comment.