Improve typing annotations for public apis

[justinchuby]

System Info

Currently when using safe_open as a context manager, mypy will complain that

"safe_open" has no attribute "__enter__"

and

"safe_open" has no attribute "__exit__"

Information

• The official example scripts
• My own modified scripts

Reproduction

Check

with safetensors.safe_open(path, framework="numpy") as f:
    ...

with a type checker.

Expected behavior

There should be no errors