Free disk size (todo)

tensorflow_datasets/core/dataset_builder.py

@@ -382,13 +382,22 @@ def _log_download_bytes(self):
     # Print is intentional: we want this to always go to stdout so user has
     # information needed to cancel download/preparation if needed.
     # This comes right before the progress bar.
+
+    dataset_size = self.info.size_in_bytes
+    def _check_disk_size():
+        if dataset_size > utils.free_disk_size(self._data_dir_root):
+            raise IOError("Not enough disk space!!\nDataset size : {dataset_size} \nFree size : {free_disk_size} \nYou need to extra {needed_disk_size} to download."
+                          .format(dataset_size=units.size_str(dataset_size),
+                                  free_disk_size=units.size_str(utils.free_disk_size(self._data_dir_root)),
+                                  needed_disk_size=units.size_str(dataset_size - utils.free_disk_size(self._data_dir_root))))
+
+    _check_disk_size()
+
     size_text = units.size_str(self.info.size_in_bytes)
     termcolor.cprint(
         "Downloading / extracting dataset %s (%s) to %s..." %
         (self.name, size_text, self._data_dir),
         attrs=["bold"])
-    # TODO(tfds): Should try to estimate the available free disk space (if
-    # possible) and raise an error if not.
 
   @abc.abstractmethod
   def _info(self):

tensorflow_datasets/core/utils/py_utils.py

@@ -28,6 +28,8 @@
 import os
 import sys
 import uuid
+import ctypes
+import platform
 
 import six
 import tensorflow as tf
@@ -275,3 +277,16 @@ def rgetattr(obj, attr, *args):
   def _getattr(obj, attr):
     return getattr(obj, attr, *args)
   return functools.reduce(_getattr, [obj] + attr.split("."))
+
+
+def free_disk_size(dirname):
+    try:
+      if platform.system() == 'Windows':
+          free_bytes = ctypes.c_ulonglong(0)
+          ctypes.windll.kernel32.GetDiskFreeSpaceExW(ctypes.c_wchar_p(dirname), None, None, ctypes.pointer(free_bytes))
+          return free_bytes.value
+      else:
+          stat = os.statvfs(dirname)
+          return stat.f_bavail * stat.f_frsize
+    except:
+      return float("inf")