IQ.Pilot Release Commit @ f2a861c

artifacts/package_sources/tinygrad/test/unit/test_tar.py

@@ -0,0 +1,141 @@
+import unittest, tarfile, io, os, pathlib, tempfile
+import numpy as np
+from tinygrad import Tensor
+from tinygrad.nn.state import tar_extract
+
+class TestTarExtractFile(unittest.TestCase):
+  def setUp(self):
+    self.test_dir = tempfile.mkdtemp()
+    self.test_files = {
+      'file1.txt': b'Hello, World!',
+      'file2.bin': b'\x00\x01\x02\x03\x04',
+      'empty_file.txt': b''
+    }
+    self.tar_path = os.path.join(self.test_dir, 'test.tar')
+    with tarfile.open(self.tar_path, 'w') as tar:
+      for filename, content in self.test_files.items():
+        file_path = os.path.join(self.test_dir, filename)
+        with open(file_path, 'wb') as f:
+          f.write(content)
+        tar.add(file_path, arcname=filename)
+
+    # Create invalid tar file
+    self.invalid_tar_path = os.path.join(self.test_dir, 'invalid.tar')
+    with open(self.invalid_tar_path, 'wb') as f:
+      f.write(b'This is not a valid tar file')
+
+  def tearDown(self):
+    for filename in self.test_files:
+      os.remove(os.path.join(self.test_dir, filename))
+    os.remove(self.tar_path)
+    os.remove(self.invalid_tar_path)
+    os.rmdir(self.test_dir)
+
+  def test_tar_extract_returns_dict(self):
+    result = tar_extract(self.tar_path)
+    self.assertIsInstance(result, dict)
+
+  def test_tar_extract_correct_keys(self):
+    result = tar_extract(self.tar_path)
+    self.assertEqual(set(result.keys()), set(self.test_files.keys()))
+
+  def test_tar_extract_content_size(self):
+    result = tar_extract(self.tar_path)
+    for filename, content in self.test_files.items():
+      self.assertEqual(len(result[filename]), len(content))
+
+  def test_tar_extract_content_values(self):
+    result = tar_extract(self.tar_path)
+    for filename, content in self.test_files.items():
+      np.testing.assert_array_equal(result[filename].numpy(), np.frombuffer(content, dtype=np.uint8))
+
+  def test_tar_extract_empty_file(self):
+    result = tar_extract(self.tar_path)
+    self.assertEqual(len(result['empty_file.txt']), 0)
+
+  def test_tar_extract_non_existent_file(self):
+    with self.assertRaises(FileNotFoundError):
+      tar_extract('non_existent_file.tar')
+
+  def test_tar_extract_invalid_file(self):
+    with self.assertRaises(tarfile.ReadError):
+      tar_extract(self.invalid_tar_path)
+
+class TestTarExtractPAX(unittest.TestCase):
+  tar_format = tarfile.PAX_FORMAT
+  max_link_len = 1000_000
+  test_files = {
+    'a/file1.txt': b'Hello, World!',
+    'a/b/file2.bin': b'\x00\x01\x02\x03\x04',
+    'empty_file.txt': b'',
+    '512file': b'a' * 512,
+    'long_file': b'some data' * 100,
+    'very' * 15 + '/' + 'very' * 15 + '_long_filename.txt': b'Hello, World!!',
+    'very' * 200 + '_long_filename.txt': b'Hello, World!!!',
+  }
+
+  def create_tar_tensor(self):
+    fobj = io.BytesIO()
+    test_dirs = set(os.path.dirname(k) for k in self.test_files.keys()).difference({ '' })
+    with tarfile.open(fileobj=fobj, mode='w', format=self.tar_format) as tar:
+      for dirname in test_dirs:
+        dir_info = tarfile.TarInfo(name=dirname)
+        dir_info.type = tarfile.DIRTYPE
+        tar.addfile(dir_info)
+
+      for filename, content in self.test_files.items():
+        file_info = tarfile.TarInfo(name=filename)
+        file_info.size = len(content)
+        tar.addfile(file_info, io.BytesIO(content))
+
+        if len(filename) < self.max_link_len:
+          link_info = tarfile.TarInfo(name=filename + '.lnk')
+          link_info.type = tarfile.SYMTYPE
+          link_info.linkname = filename
+          tar.addfile(link_info)
+    return Tensor(fobj.getvalue())
+
+  def test_tar_extract_returns_dict(self):
+    result = tar_extract(self.create_tar_tensor())
+    self.assertIsInstance(result, dict)
+
+  def test_tar_extract_correct_keys(self):
+    result = tar_extract(self.create_tar_tensor())
+    self.assertEqual(set(result.keys()), set(self.test_files.keys()))
+
+  def test_tar_extract_content_size(self):
+    result = tar_extract(self.create_tar_tensor())
+    for filename, content in self.test_files.items():
+      self.assertEqual(len(result[filename]), len(content))
+
+  def test_tar_extract_content_values(self):
+    result = tar_extract(self.create_tar_tensor())
+    for filename, content in self.test_files.items():
+      np.testing.assert_array_equal(result[filename].numpy(), np.frombuffer(content, dtype=np.uint8))
+
+  def test_tar_extract_empty_file(self):
+    result = tar_extract(self.create_tar_tensor())
+    self.assertEqual(len(result['empty_file.txt']), 0)
+
+  def test_tar_extract_non_existent_file(self):
+    with self.assertRaises(FileNotFoundError):
+      tar_extract(Tensor(pathlib.Path('non_existent_file.tar')))
+
+  def test_tar_extract_invalid_file(self):
+    with self.assertRaises(tarfile.ReadError):
+      tar_extract(Tensor(b'This is not a valid tar file'))
+
+  def test_tar_extract_invalid_file_long(self):
+    with self.assertRaises(tarfile.ReadError):
+      tar_extract(Tensor(b'This is not a valid tar file'*100))
+
+class TestTarExtractUSTAR(TestTarExtractPAX):
+  tar_format = tarfile.USTAR_FORMAT
+  max_link_len = 100
+  test_files = {k: v for k, v in TestTarExtractPAX.test_files.items() if len(k) < 256}
+
+class TestTarExtractGNU(TestTarExtractPAX):
+  tar_format = tarfile.GNU_FORMAT
+
+if __name__ == '__main__':
+  unittest.main()