fix: address review comments (iteration #2)

Author: aviruthen

sagemaker-core/src/sagemaker/core/utils/__init__.py

@@ -19,9 +19,11 @@
 """
 from __future__ import absolute_import
 
-# Public API surface: only non-private functions are exported via __all__.
+# Public API surface.
+# Note: _save_model is underscore-prefixed but was already in __all__ (pre-existing).
+# custom_extractall_tarfile is the main public entry point for safe tar extraction.
 # Private helpers (_get_resolved_path, _is_bad_path, _is_bad_link, _get_safe_members)
-# are still importable directly but are not part of the public API.
+# are importable directly from sagemaker.core.common_utils but are not re-exported here.
 __all__ = [
     "_save_model",
     "download_file_from_url",
@@ -41,18 +43,10 @@
     "get_config_value",
 ]
 
-# Internal helpers that are importable but not part of the public API
-_INTERNAL_NAMES = [
-    "_get_resolved_path",
-    "_is_bad_path",
-    "_is_bad_link",
-    "_get_safe_members",
-]
-
 
 def __getattr__(name):
     """Lazy import to avoid circular dependencies."""
-    if name in __all__ or name in _INTERNAL_NAMES:
+    if name in __all__:
         from sagemaker.core import common_utils
 
         return getattr(common_utils, name)

sagemaker-core/tests/unit/test_common_utils_tar_safety.py

@@ -15,10 +15,8 @@
 
 import os
 import tempfile
-import tarfile
 
-import pytest
-from unittest.mock import Mock, patch, MagicMock
+from unittest.mock import Mock, patch
 
 from sagemaker.core.common_utils import (
     _get_resolved_path,
@@ -204,18 +202,20 @@ def test_get_safe_members_filters_bad_hardlink_member():
 def test_custom_extractall_tarfile_with_data_filter_uses_filter_param():
     """Test custom_extractall_tarfile uses data_filter when available.
 
-    We set mock_tarfile.data_filter explicitly to ensure hasattr returns True.
-    The MagicMock would auto-create the attribute anyway, but we set it
-    explicitly for clarity. The key assertion is that filter="data" is passed.
+    We patch the module-level `tarfile` import in common_utils (not the `tar` parameter).
+    Setting mock_tarfile.data_filter explicitly ensures hasattr(tarfile, 'data_filter')
+    returns True inside custom_extractall_tarfile. The key assertion is that
+    filter="data" is passed to tar.extractall.
     """
     mock_tar = Mock()
     mock_tar.extractall = Mock()
 
     with tempfile.TemporaryDirectory() as tmpdir:
         extract_path = os.path.join(tmpdir, "extract")
 
+        # Patch the module-level tarfile import in common_utils
         with patch("sagemaker.core.common_utils.tarfile") as mock_tarfile:
-            # Explicitly set data_filter to ensure the hasattr check passes
+            # Explicitly set data_filter so hasattr check passes
             mock_tarfile.data_filter = True
 
             custom_extractall_tarfile(mock_tar, extract_path)
@@ -226,6 +226,9 @@ def test_custom_extractall_tarfile_with_data_filter_uses_filter_param():
 def test_custom_extractall_tarfile_without_data_filter_uses_safe_members():
     """Test custom_extractall_tarfile uses safe members when data_filter is unavailable.
 
+    Uses spec=['TarFile'] to restrict the mock so that
+    hasattr(mock_tarfile, 'data_filter') returns False, forcing the fallback path.
+
     Verifies that:
     1. tar.getmembers() is called (not iterating over tar directly)
     2. _get_safe_members is called with the members list and resolved extract_path as base