import sys
from pylease.logger import LOGME as logme  # noqa
from pylease.command.task import AfterTask
from pylease.extension import Extension


class PypiExtension(Extension):
    def load(self):
        make_command = self._lizy.commands['make']
        parser = make_command.parser

        pypi_group = parser.add_argument_group(title='PyPI arguments')

        pypi_group.add_argument('--pypi',
                                action='store_true',
                                dest='to_pypi',
                                default=False,
                                help='Upload package to PyPI')

        make_command.add_after_task(PypiAfterTask(rollback))


class PypiAfterTask(AfterTask):  # pragma: no cover - Unable to test this other than manually TODO: try to test
    def execute(self, lizy, args):
        if args.to_pypi:
            self.enable_rollback()

            orig_args = sys.argv
            sys.argv = ['setup.py', 'sdist', 'upload']
            __import__('setup')
            sys.argv = orig_args


def rollback():  # pragma: no cover
    logme.debug('Unable to rollback PyPI upload, please delete the uploaded version manually.')

@@ -4,16 +4,69 @@
 ort sys%0A
+from pylease.command.rollback import Rollback, Stage%0A
 from pyl
@@ -709,24 +709,16 @@
 terTask(
-rollback
 ))%0A%0A%0Acla
@@ -923,43 +923,29 @@
 ack(
-)%0A%0A            orig_args = 
+PypiRollback(
 sys.argv
 %0A   
@@ -940,16 +940,19 @@
 sys.argv
+))%0A
 %0A       
@@ -1015,48 +1015,391 @@
     
-__import__('setup')%0A            sys.argv
+self._rollback.enable_stage(PypiRollback.UPLOADING_STAGE)%0A%0A            __import__('setup')%0A            self._rollback.enable_stage(PypiRollback.UPLOADED_STAGE)%0A%0A%0Aclass PypiRollback(Rollback):  # pragma: no cover%0A    UPLOADED_STAGE = 'uploaded'%0A    UPLOADING_STAGE = 'uploading'%0A%0A    def __init__(self, orig_args):%0A        super(PypiRollback, self).__init__()%0A%0A        self._orig_args
  = o
@@ -1408,17 +1408,48 @@
 g_args%0A%0A
-%0A
+    @Stage(UPLOADING_STAGE)%0A    
 def roll
@@ -1456,32 +1456,119 @@
 back
-(
+_args(self
 ):
+%0A
   
-# pragma: no cover%0A
+      sys.argv = self._orig_args%0A%0A    @Stage(UPLOADED_STAGE)%0A    def rollback_inform(self):%0A    
     
@@ -1577,13 +1577,12 @@
 gme.
-debug
+warn
 ('Un

Improve PyPI extension rollbacks