from setuptools import setup
from setuptools.command.install import install
import shutil
import os
import atexit

class CustomInstallCommand(install):
     """Customized install command to perform cleanup after setup."""

     def run(self):
          # Run the standard install
          super().run()

          # Register cleanup function to run at the end of the setup process
          atexit.register(self.cleanup)

     def cleanup(self):
          """Cleanup directories generated by build."""
          for f in ['build', 'idapro.egg-info']:
               if os.path.exists(f):
                    shutil.rmtree(f)

setup(
     name="idapro",
     version="0.0.2",
     python_requires=">=3.8",
     packages=["idapro"],
     cmdclass={
          'install': CustomInstallCommand,
     },
)