Fix version overide

Author: vijayphoenix

tensorflow_datasets/core/registered.py

@@ -109,7 +109,7 @@ def __init__(self, name, is_abstract=False, in_development=False):
 class RegisteredDataset(abc.ABCMeta):
   """Subclasses will be registered and given a `name` property."""
 
-  def __new__(cls, cls_name, bases, class_dict): # pylint: disable= C0204, W0221
+  def __new__(cls, cls_name, bases, class_dict): # pylint: disable = arguments-differ, bad-mcs-classmethod-argument
     name = naming.camelcase_to_snakecase(cls_name)
     class_dict["name"] = name
     builder_cls = super(RegisteredDataset, cls).__new__(  # pylint: disable=too-many-function-args
@@ -164,8 +164,7 @@ def builder(name, **builder_init_kwargs):
     DatasetNotFoundError: if `name` is unrecognized.
   """
   name, builder_kwargs = _dataset_name_and_kwargs_from_name_str(name)
-  builder_kwargs.update( \
-    (k, v) for k, v in builder_init_kwargs.items() if v is not None)
+  builder_kwargs.update(builder_init_kwargs)
   if name in _ABSTRACT_DATASET_REGISTRY:
     raise DatasetNotFoundError(name, is_abstract=True)
   if name in _IN_DEVELOPMENT_REGISTRY:
@@ -294,8 +293,7 @@ def load(name,
   # pylint: enable=line-too-long
 
   name, name_builder_kwargs = _dataset_name_and_kwargs_from_name_str(name)
-  name_builder_kwargs.update( \
-    (k, v) for k, v in (builder_kwargs or {}).items() if v is not None)
+  name_builder_kwargs.update(builder_kwargs or {})
   builder_kwargs = name_builder_kwargs
 
   # Set data_dir

tensorflow_datasets/scripts/download_and_prepare.py

@@ -180,8 +180,11 @@ def main(_):
   logging.info("Running download_and_prepare for datasets:\n%s",
                "\n".join(datasets_to_build))
   logging.info('Version: "%s"', version)
+  kwargs = {"data_dir": FLAGS.data_dir}
+  if version is not None:
+    kwargs.update(("version", version))
   builders = {
-      name: tfds.builder(name, data_dir=FLAGS.data_dir, version=version)
+      name: tfds.builder(name, **kwargs)
       for name in datasets_to_build
   }
 
@@ -196,18 +199,17 @@ def main(_):
           "--builder_config_id can only be used with datasets with configs")
     config = builder.BUILDER_CONFIGS[FLAGS.builder_config_id]
     logging.info("Running download_and_prepare for config: %s", config.name)
-    builder_for_config = tfds.builder(
-        builder.name, data_dir=FLAGS.data_dir, config=config, version=version)
+    kwargs.update({("config", config)})
+    builder_for_config = tfds.builder(builder.name, **kwargs)
     download_and_prepare(builder_for_config)
   else:
     for name, builder in builders.items():
       if builder.BUILDER_CONFIGS and "/" not in name:
         # If builder has multiple configs, and no particular config was
         # requested, then compute all.
         for config in builder.BUILDER_CONFIGS:
-          builder_for_config = tfds.builder(
-              builder.name, data_dir=FLAGS.data_dir, config=config,
-              version=version)
+          kwargs.update({("config", config)})
+          builder_for_config = tfds.builder(name, **kwargs)
           download_and_prepare(builder_for_config)
       else:
         # If there is a slash in the name, then user requested a specific