Before writing custom code for this issue, take a look at the step export/import settings. NX (and I assume other CAD systems as well) has the ability to save the assembly structure in the step file. When opened in NX, it will create a new part file for each original component. You should be able to work with the file originator to find export settings that work best for you.
Normally, the assembly file does what you say but these step files are like single part files with multiple "bodies" (the correct technical word is feature?)