diff --git a/docker_squash/cli.py b/docker_squash/cli.py
index 462b872..cff0773 100644
--- a/docker_squash/cli.py
+++ b/docker_squash/cli.py
@@ -42,7 +42,9 @@ class CLI(object):
         parser.add_argument(
             '-f', '--from-layer', help='ID of the layer or image ID or image name. If not specified will squash all layers in the image')
         parser.add_argument(
-            '-t', '--tag', help="Specify the tag to be used for the new image. By default it'll be set to 'image' argument")
+            '-t', '--tag', help="Specify the tag to be used for the new image. by default it'll be set to 'image' argument")
+        parser.add_argument(
+            '-c', '--cleanup', action='store_true', help="Remove source image from Docker after squashing")
         parser.add_argument(
             '--tmp-dir', help='Temporary directory to be used')
         parser.add_argument(
@@ -59,7 +61,7 @@ class CLI(object):
 
         try:
             squash.Squash(log=self.log, image=args.image,
-                          from_layer=args.from_layer, tag=args.tag, output_path=args.output_path, tmp_dir=args.tmp_dir, development=args.development).run()
+                          from_layer=args.from_layer, tag=args.tag, output_path=args.output_path, tmp_dir=args.tmp_dir, development=args.development, cleanup=args.cleanup).run()
         except KeyboardInterrupt:
             self.log.error("Program interrupted by user, exiting...")
             sys.exit(1)
diff --git a/docker_squash/squash.py b/docker_squash/squash.py
index d6d21e0..e81960a 100644
--- a/docker_squash/squash.py
+++ b/docker_squash/squash.py
@@ -14,7 +14,7 @@ from docker_squash.version import version
 class Squash(object):
 
     def __init__(self, log, image, docker=None, from_layer=None, tag=None, tmp_dir=None,
-                 output_path=None, load_image=True, development=False):
+                 output_path=None, load_image=True, development=False, cleanup=False):
         self.log = log
         self.docker = docker
         self.image = image
@@ -24,6 +24,7 @@ class Squash(object):
         self.output_path = output_path
         self.load_image = load_image
         self.development = development
+        self.cleanup = cleanup
 
         if not docker:
             self.docker = common.docker_client(self.log)
@@ -65,12 +66,22 @@ class Squash(object):
 
             raise
 
+    def _cleanup(self):
+        try:
+            image_id = self.docker.inspect_image(self.image)['Id']
+        except:
+            self.log.warn("Could not get the image ID for %s image, skipping cleanup after squashing" % self.image)
+            return
+
+        self.log.info("Removing old %s image..." % self.image)
+        self.docker.remove_image(image_id, force=False, noprune=False)
+        self.log.info("Image removed!")
+
     def squash(self, image):
         # Do the actual squashing
         new_image_id = image.squash()
 
-        self.log.info(
-            "New squashed image ID is %s" % new_image_id)
+        self.log.info("New squashed image ID is %s" % new_image_id)
 
         if self.output_path:
             # Move the tar archive to the specified path
@@ -80,8 +91,16 @@ class Squash(object):
             # Load squashed image into Docker
             image.load_squashed_image()
 
+        # Clean up all temporary files
         image.cleanup()
 
+        # Remove the source image - this is the only possible way
+        # to remove orphaned layers from Docker daemon at the build time.
+        # We cannot use here a tag name because it could be used as the target,
+        # squashed image tag - we need to use the image ID.
+        if self.cleanup:
+            self._cleanup()
+
         self.log.info("Done")
 
         return new_image_id
diff --git a/docker_squash/v2_image.py b/docker_squash/v2_image.py
index c2fa1ac..4df00d7 100644
--- a/docker_squash/v2_image.py
+++ b/docker_squash/v2_image.py
@@ -280,10 +280,10 @@ class V2Image(Image):
         if self.layer_paths_to_move:
             config['parent'] = self.layer_paths_to_move[-1]
         else:
-            del config['parent']
+            config.pop("parent", None)
         # Update 'id' - it should be the path to the layer
         config['id'] = layer_path_id
-        del config['container']
+        config.pop("container", None)
         return config
 
     def _generate_image_metadata(self):
@@ -294,7 +294,7 @@ class V2Image(Image):
         metadata['created'] = self.date
 
         # Remove unnecessary or old fields
-        del metadata['container']
+        metadata.pop("container", None)
 
         # Remove squashed layers from history
         metadata['history'] = metadata['history'][:len(self.layers_to_move)]

Remove squashed layers from Docker daemon
In most cases we're not interested in layers that were squashed. In case where the input and target image has the same tag - all squashed layers will be "orphaned" leaving them in Docker daemon and wasting space. We should probably include a `--cleanup` switch or something similar to remove layers without tag that were "orphaned" during the squashing process.
diff --git a/tests/test_unit_squash.py b/tests/test_unit_squash.py
new file mode 100644
index 0000000..c136761
--- /dev/null
+++ b/tests/test_unit_squash.py
@@ -0,0 +1,43 @@
+import unittest
+import mock
+import six
+import logging
+
+from docker_squash.squash import Squash
+from docker_squash.image import Image
+from docker_squash.errors import SquashError
+
+
+class TestSquash(unittest.TestCase):
+    
+    def setUp(self):
+        self.log = mock.Mock()
+        self.docker_client = mock.Mock()
+        self.docker_client.version.return_value = {'GitCommit': "commit/9.9.9", 'ApiVersion': "9.99"}
+
+    def test_handle_case_when_no_image_is_provided(self):
+        squash = Squash(self.log, None, self.docker_client)
+        with self.assertRaises(SquashError) as cm:
+            squash.run()
+        self.assertEquals(
+            str(cm.exception), "Image is not provided")
+
+    def test_exit_if_no_output_path_provided_and_loading_is_disabled_too(self):
+        squash = Squash(self.log, 'image', self.docker_client, load_image=False, output_path=None)
+        squash.run()
+        self.log.warn.assert_called_with("No output path specified and loading into Docker is not selected either; squashed image would not accessible, proceeding with squashing doesn't make sense")
+
+    @mock.patch('docker_squash.squash.V2Image')
+    def test_should_not_cleanup_after_squashing(self, v2_image):
+        squash = Squash(self.log, 'image', self.docker_client, load_image=True)
+        squash.run()
+    
+        v2_image.cleanup.assert_not_called()
+
+    @mock.patch('docker_squash.squash.V2Image')
+    def test_should_cleanup_after_squashing(self, v2_image):
+        self.docker_client.inspect_image.return_value = {'Id': "abcdefgh"}
+        squash = Squash(self.log, 'image', self.docker_client, load_image=True, cleanup=True)
+        squash.run()
+    
+        self.docker_client.remove_image.assert_called_with('abcdefgh', force=False, noprune=False)
diff --git a/tests/test_unit_v1_image.py b/tests/test_unit_v1_image.py
index 6316399..07081ab 100644
--- a/tests/test_unit_v1_image.py
+++ b/tests/test_unit_v1_image.py
@@ -299,24 +299,5 @@ class TestAddMarkers(unittest.TestCase):
         self.assertTrue(tar_info.isfile())
 
 
-class TestGeneral(unittest.TestCase):
-
-    def setUp(self):
-        self.log = mock.Mock()
-        self.docker_client = mock.Mock()
-        self.docker_client.version.return_value = {'GitCommit': "commit/9.9.9", 'ApiVersion': "9.99"}
-
-    def test_handle_case_when_no_image_is_provided(self):
-        squash = Squash(self.log, None, self.docker_client)
-        with self.assertRaises(SquashError) as cm:
-            squash.run()
-        self.assertEquals(
-            str(cm.exception), "Image is not provided")
-
-    def test_exit_if_no_output_path_provided_and_loading_is_disabled_too(self):
-        squash = Squash(self.log, 'image', self.docker_client, load_image=False, output_path=None)
-        squash.run()
-        self.log.warn.assert_called_with("No output path specified and loading into Docker is not selected either; squashed image would not accessible, proceeding with squashing doesn't make sense")
-
 if __name__ == '__main__':
     unittest.main()

pip install -e .[dev]
parse_log_pytest
requirements.txt
pytest pytest-cov mock
apt-get update apt-get install -y gcc
3.9
requirements.txt
pytest --no-header -rA --tb=line --color=no -p no:cacheprovider -W ignore::DeprecationWarning
certifi==2025.1.31
charset-normalizer==3.4.1
coverage==7.8.0
docker-py==1.10.6
docker-pycreds==0.4.0
-e git+https://github.com/goldmann/docker-squash.git@15bda09f82d4031fcf3e493d18133ffbbd09fdbd#egg=docker_squash
exceptiongroup==1.2.2
idna==3.10
iniconfig==2.1.0
mock==5.2.0
packaging==24.2
pluggy==1.5.0
pytest==8.3.5
pytest-cov==6.0.0
requests==2.32.3
six==1.17.0
tomli==2.2.1
urllib3==2.3.0
websocket-client==1.8.0

name: docker-squash
channels:
  - defaults
  - https://repo.anaconda.com/pkgs/main
  - https://repo.anaconda.com/pkgs/r
  - conda-forge
dependencies:
  - _libgcc_mutex=0.1=main
  - _openmp_mutex=5.1=1_gnu
  - ca-certificates=2025.2.25=h06a4308_0
  - ld_impl_linux-64=2.40=h12ee557_0
  - libffi=3.4.4=h6a678d5_1
  - libgcc-ng=11.2.0=h1234567_1
  - libgomp=11.2.0=h1234567_1
  - libstdcxx-ng=11.2.0=h1234567_1
  - ncurses=6.4=h6a678d5_0
  - openssl=3.0.16=h5eee18b_0
  - pip=25.0=py39h06a4308_0
  - python=3.9.21=he870216_1
  - readline=8.2=h5eee18b_0
  - setuptools=75.8.0=py39h06a4308_0
  - sqlite=3.45.3=h5eee18b_0
  - tk=8.6.14=h39e8969_0
  - tzdata=2025a=h04d1e81_0
  - wheel=0.45.1=py39h06a4308_0
  - xz=5.6.4=h5eee18b_1
  - zlib=1.2.13=h5eee18b_1
  - pip:
      - certifi==2025.1.31
      - charset-normalizer==3.4.1
      - coverage==7.8.0
      - docker-py==1.10.6
      - docker-pycreds==0.4.0
      - exceptiongroup==1.2.2
      - idna==3.10
      - iniconfig==2.1.0
      - mock==5.2.0
      - packaging==24.2
      - pluggy==1.5.0
      - pytest==8.3.5
      - pytest-cov==6.0.0
      - requests==2.32.3
      - six==1.17.0
      - tomli==2.2.1
      - urllib3==2.3.0
      - websocket-client==1.8.0
prefix: /opt/conda/envs/docker-squash

tests/test_unit_squash.py::TestSquash::test_should_cleanup_after_squashing
tests/test_unit_squash.py::TestSquash::test_exit_if_no_output_path_provided_and_loading_is_disabled_too tests/test_unit_squash.py::TestSquash::test_handle_case_when_no_image_is_provided tests/test_unit_squash.py::TestSquash::test_should_not_cleanup_after_squashing tests/test_unit_v1_image.py::TestSkippingFiles::test_should_not_skip_file_not_in_path_to_skip tests/test_unit_v1_image.py::TestSkippingFiles::test_should_not_skip_the_file_that_name_is_similar_to_skipped_path tests/test_unit_v1_image.py::TestSkippingFiles::test_should_skip_exact_files tests/test_unit_v1_image.py::TestSkippingFiles::test_should_skip_files_in_subdirectory tests/test_unit_v1_image.py::TestParseImageName::test_should_parse_name_name_with_proper_tag tests/test_unit_v1_image.py::TestParseImageName::test_should_parse_name_name_without_tag tests/test_unit_v1_image.py::TestPrepareTemporaryDirectory::test_create_tmp_directory_if_not_provided tests/test_unit_v1_image.py::TestPrepareTemporaryDirectory::test_should_raise_if_directory_already_exists tests/test_unit_v1_image.py::TestPrepareTemporaryDirectory::test_should_use_provided_tmp_dir tests/test_unit_v1_image.py::TestPrepareLayersToSquash::test_should_generate_list_of_layers tests/test_unit_v1_image.py::TestPrepareLayersToSquash::test_should_not_fail_with_empty_list_of_layers tests/test_unit_v1_image.py::TestPrepareLayersToSquash::test_should_return_all_layers_if_from_layer_is_not_found tests/test_unit_v1_image.py::TestGenerateV1ImageId::test_should_generate_id tests/test_unit_v1_image.py::TestGenerateV1ImageId::test_should_generate_id_that_is_not_integer_shen_shortened tests/test_unit_v1_image.py::TestGenerateRepositoriesJSON::test_generate_json tests/test_unit_v1_image.py::TestGenerateRepositoriesJSON::test_handle_empty_image_id tests/test_unit_v1_image.py::TestGenerateRepositoriesJSON::test_should_not_generate_repositories_if_name_and_tag_is_missing tests/test_unit_v1_image.py::TestMarkerFiles::test_should_find_all_marker_files tests/test_unit_v1_image.py::TestMarkerFiles::test_should_return_empty_dict_when_no_files_are_in_the_tar tests/test_unit_v1_image.py::TestMarkerFiles::test_should_return_empty_dict_when_no_marker_files_are_found tests/test_unit_v1_image.py::TestAddMarkers::test_should_add_all_marker_files_to_empty_tar tests/test_unit_v1_image.py::TestAddMarkers::test_should_add_marker_file_when_tar_has_prefixed_entries tests/test_unit_v1_image.py::TestAddMarkers::test_should_not_add_any_marker_files tests/test_unit_v1_image.py::TestAddMarkers::test_should_not_fail_with_empty_list_of_markers_to_add tests/test_unit_v1_image.py::TestAddMarkers::test_should_skip_a_marker_file_if_file_is_in_unsquashed_layers
docker_squash/v2_image.py docker_squash/squash.py docker_squash/cli.py
docker_squash/v2_image.py docker_squash/squash.py docker_squash/cli.py