Repository: IDSIA/sacred
Pull request #397 (2019-01-02 13:12:16)

Attach content-type directly to gridfs object
GridFS offers a special attribute to attach a content-type to a file.
This commit introduces an extra arugment to add_artifact to put the
content-type into. If the content-type is not set manually it is
automatically detected. However, this behavior makes more sense then
reserving a field in the artifacts metadata for the content-type. This
change has only an effect on the MongoObserver.


Patch:
diff --git a/sacred/experiment.py b/sacred/experiment.py
index bd93f56..47d5199 100755
--- a/sacred/experiment.py
+++ b/sacred/experiment.py
@@ -339,7 +339,7 @@ class Experiment(Ingredient):
         assert self.current_run is not None, "Can only be called during a run."
         self.current_run.add_resource(filename)
 
-    def add_artifact(self, filename, name=None, metadata=None):
+    def add_artifact(self, filename, name=None, metadata=None, content_type=None):
         """Add a file as an artifact.
 
         In Sacred terminology an artifact is a file produced by the experiment
@@ -359,9 +359,12 @@ class Experiment(Ingredient):
         metadata: dict, optional
             optionally attach metadata to the artifact.
             This only has an effect when using the MongoObserver.
+        content_type: str, optional
+            optionally attach a content-type to the artifact.
+            This only has an effect when using the MongoObserver.
         """
         assert self.current_run is not None, "Can only be called during a run."
-        self.current_run.add_artifact(filename, name, metadata)
+        self.current_run.add_artifact(filename, name, metadata, content_type)
 
     @property
     def info(self):
diff --git a/sacred/observers/base.py b/sacred/observers/base.py
index a41a42c..1eceed6 100644
--- a/sacred/observers/base.py
+++ b/sacred/observers/base.py
@@ -33,5 +33,5 @@ class RunObserver(object):
     def resource_event(self, filename):
         pass
 
-    def artifact_event(self, name, filename, metadata=None):
+    def artifact_event(self, name, filename, metadata=None, content_type=None):
         pass
diff --git a/sacred/observers/file_storage.py b/sacred/observers/file_storage.py
index 7c1f793..0b66726 100644
--- a/sacred/observers/file_storage.py
+++ b/sacred/observers/file_storage.py
@@ -206,7 +206,7 @@ class FileStorageObserver(RunObserver):
         self.run_entry['resources'].append([filename, store_path])
         self.save_json(self.run_entry, 'run.json')
 
-    def artifact_event(self, name, filename, metadata=None):
+    def artifact_event(self, name, filename, metadata=None, content_type=None):
         self.save_file(filename, name)
         self.run_entry['artifacts'].append(name)
         self.save_json(self.run_entry, 'run.json')
diff --git a/sacred/observers/mongo.py b/sacred/observers/mongo.py
index faed914..6cfc6c5 100644
--- a/sacred/observers/mongo.py
+++ b/sacred/observers/mongo.py
@@ -187,36 +187,30 @@ class MongoObserver(RunObserver):
         self.run_entry['resources'].append((filename, md5hash))
         self.save()
 
-    def artifact_event(self, name, filename, metadata=None):
+    def artifact_event(self, name, filename, metadata=None, content_type=None):
         with open(filename, 'rb') as f:
             run_id = self.run_entry['_id']
             db_filename = 'artifact://{}/{}/{}'.format(self.runs.name, run_id,
                                                        name)
-            if not self._content_type_manually_set(metadata):
-                metadata = self._try_to_detect_content_type(metadata, filename)
+            if content_type is None:
+                content_type = self._try_to_detect_content_type(filename)
 
-            file_id = self.fs.put(f, filename=db_filename, metadata=metadata)
+            file_id = self.fs.put(f, filename=db_filename, metadata=metadata, content_type=content_type)
 
         self.run_entry['artifacts'].append({'name': name,
                                             'file_id': file_id})
         self.save()
 
-    def _content_type_manually_set(self, metadata):
-        return metadata is not None and 'content-type' in metadata
-
-    def _try_to_detect_content_type(self, metadata, filename):
+    @staticmethod
+    def _try_to_detect_content_type(filename):
         mime_type, _ = mimetypes.guess_type(filename)
         if mime_type is not None:
-            if metadata is None:
-                metadata = {}
-            content_type_metadata = {'content-type': mime_type}
-            metadata.update(content_type_metadata)
-            print('Added {} to metadata of artifact {}.'.format(
-                content_type_metadata, filename))
+            print('Added {} as content-type of artifact {}.'.format(
+                mime_type, filename))
         else:
             print('Failed to detect content-type automatically for '
                   'artifact {}.'.format(filename))
-        return metadata
+        return mime_type
 
     def log_metrics(self, metrics_by_name, info):
         """Store new measurements to the database.
diff --git a/sacred/observers/sql.py b/sacred/observers/sql.py
index a84b47b..ee0daa0 100644
--- a/sacred/observers/sql.py
+++ b/sacred/observers/sql.py
@@ -106,7 +106,7 @@ class SqlObserver(RunObserver):
         self.run.resources.append(res)
         self.save()
 
-    def artifact_event(self, name, filename, metadata=None):
+    def artifact_event(self, name, filename, metadata=None, content_type=None):
         a = Artifact.create(name, filename)
         self.run.artifacts.append(a)
         self.save()
diff --git a/sacred/observers/tinydb_hashfs.py b/sacred/observers/tinydb_hashfs.py
index 1a17cf4..9c272c4 100644
--- a/sacred/observers/tinydb_hashfs.py
+++ b/sacred/observers/tinydb_hashfs.py
@@ -189,7 +189,7 @@ class TinyDbObserver(RunObserver):
             self.run_entry['resources'].append(resource)
             self.save()
 
-    def artifact_event(self, name, filename, metadata=None):
+    def artifact_event(self, name, filename, metadata=None, content_type=None):
 
         id_ = self.fs.put(filename).id
         handle = BufferedReaderWrapper(open(filename, 'rb'))
diff --git a/sacred/run.py b/sacred/run.py
index 7848d50..18caf66 100755
--- a/sacred/run.py
+++ b/sacred/run.py
@@ -157,7 +157,7 @@ class Run(object):
         filename = os.path.abspath(filename)
         self._emit_resource_added(filename)
 
-    def add_artifact(self, filename, name=None, metadata=None):
+    def add_artifact(self, filename, name=None, metadata=None, content_type=None):
         """Add a file as an artifact.
 
         In Sacred terminology an artifact is a file produced by the experiment
@@ -176,10 +176,13 @@ class Run(object):
         metadata: dict
             optionally attach metadata to the artifact.
             This only has an effect when using the MongoObserver.
+        content_type: str, optional
+            optionally attach a content-type to the artifact.
+            This only has an effect when using the MongoObserver.
         """
         filename = os.path.abspath(filename)
         name = os.path.basename(filename) if name is None else name
-        self._emit_artifact_added(name, filename, metadata)
+        self._emit_artifact_added(name, filename, metadata, content_type)
 
     def __call__(self, *args):
         r"""Start this run.
@@ -377,12 +380,13 @@ class Run(object):
         for observer in self.observers:
             self._safe_call(observer, 'resource_event', filename=filename)
 
-    def _emit_artifact_added(self, name, filename, metadata):
+    def _emit_artifact_added(self, name, filename, metadata, content_type):
         for observer in self.observers:
             self._safe_call(observer, 'artifact_event',
                             name=name,
                             filename=filename,
-                            metadata=metadata)
+                            metadata=metadata,
+                            content_type=content_type)
 
     def _safe_call(self, obs, method, **kwargs):
         if obs not in self._failed_observers and hasattr(obs, method):


Test patch:
diff --git a/tests/test_observers/test_mongo_observer.py b/tests/test_observers/test_mongo_observer.py
index db34ae7..9500d71 100644
--- a/tests/test_observers/test_mongo_observer.py
+++ b/tests/test_observers/test_mongo_observer.py
@@ -323,8 +323,8 @@ def test_log_metrics(mongo_obs, sample_run, logged_metrics):
     assert mongo_obs.metrics.count() == 4
 
 
-def test_mongo_observer_artifact_event_content_type_detected(mongo_obs, sample_run):
-    """Test that the content-type is detected for artifacts."""
+def test_mongo_observer_artifact_event_content_type_added(mongo_obs, sample_run):
+    """Test that the detected content_type is added to other metadata."""
     mongo_obs.started_event(**sample_run)
 
     filename = 'setup.py'
@@ -333,42 +333,39 @@ def test_mongo_observer_artifact_event_content_type_detected(mongo_obs, sample_r
     mongo_obs.artifact_event(name, filename)
 
     assert mongo_obs.fs.put.called
-    assert mongo_obs.fs.put.call_args[1]['metadata']['content-type'] == 'text/x-python'
+    assert mongo_obs.fs.put.call_args[1]['content_type'] == 'text/x-python'
 
     db_run = mongo_obs.runs.find_one()
     assert db_run['artifacts']
 
 
-def test_mongo_observer_artifact_event_content_type_added(mongo_obs, sample_run):
-    """Test that the detected content-type is added to other metadata."""
+def test_mongo_observer_artifact_event_content_type_not_overwritten(mongo_obs, sample_run):
+    """Test that manually set content_type is not overwritten by automatic detection."""
     mongo_obs.started_event(**sample_run)
 
     filename = 'setup.py'
     name = 'mysetup'
 
-    mongo_obs.artifact_event(name, filename, metadata={'comment': 'the setup file'})
+    mongo_obs.artifact_event(name, filename, content_type='application/json')
 
     assert mongo_obs.fs.put.called
-    assert mongo_obs.fs.put.call_args[1]['metadata']['content-type'] == 'text/x-python'
-    assert mongo_obs.fs.put.call_args[1]['metadata']['comment'] == 'the setup file'
+    assert mongo_obs.fs.put.call_args[1]['content_type'] == 'application/json'
 
     db_run = mongo_obs.runs.find_one()
     assert db_run['artifacts']
 
 
-def test_mongo_observer_artifact_event_content_type_not_overwritten(mongo_obs, sample_run):
-    """Test that manually set content-type
-    metadata is not overwritten by automatic detection.
-    """
+def test_mongo_observer_artifact_event_metadata(mongo_obs, sample_run):
+    """Test that the detected content-type is added to other metadata."""
     mongo_obs.started_event(**sample_run)
 
     filename = 'setup.py'
     name = 'mysetup'
 
-    mongo_obs.artifact_event(name, filename, metadata={'content-type': 'application/json'})
+    mongo_obs.artifact_event(name, filename, metadata={'comment': 'the setup file'})
 
     assert mongo_obs.fs.put.called
-    assert mongo_obs.fs.put.call_args[1]['metadata']['content-type'] == 'application/json'
+    assert mongo_obs.fs.put.call_args[1]['metadata']['comment'] == 'the setup file'
 
     db_run = mongo_obs.runs.find_one()
     assert db_run['artifacts']
diff --git a/tests/test_run.py b/tests/test_run.py
index f732689..700a643 100755
--- a/tests/test_run.py
+++ b/tests/test_run.py
@@ -129,9 +129,12 @@ def test_run_heartbeat_event(run):
 def test_run_artifact_event(run):
     observer = run.observers[0]
     handle, f_name = tempfile.mkstemp()
+    name = 'foobar'
     metadata = {'testkey': 42}
-    run.add_artifact(f_name, name='foobar', metadata=metadata)
-    observer.artifact_event.assert_called_with(filename=f_name, name='foobar', metadata=metadata)
+    content_type = 'text/plain'
+    run.add_artifact(f_name, name=name, metadata=metadata, content_type=content_type)
+    observer.artifact_event.assert_called_with(filename=f_name, name=name,
+                                               metadata=metadata, content_type=content_type)
     os.close(handle)
     os.remove(f_name)