[DataPipe] Improve .pyi generation (#72829)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/72829
Make two functions more flexible and usable from a different repo.
Test Plan: Imported from OSS
Reviewed By: ejguan
Differential Revision: D34227912
Pulled By: NivekT
fbshipit-source-id: 873934ed33caf485de7f56e9c4a1d3f3fa1a92ef
(cherry picked from commit b990c5e4c7244e1d0352b24c475df2c0968ee1c0)
