Repository: iterative/cml
Pull request #337 (2020-11-06 11:38:04)

Release v0.2.2 (#331 #335)
#331 
#335 


Patch:
diff --git a/bin/cml-publish.js b/bin/cml-publish.js
index c211003..a86cc86 100644
--- a/bin/cml-publish.js
+++ b/bin/cml-publish.js
@@ -10,14 +10,25 @@ const yargs = require('yargs');
 const CML = require('../src/cml');
 
 const run = async (opts) => {
-  const { data, file, 'gitlab-uploads': gitlab_uploads } = opts;
+  const {
+    data,
+    file,
+    'gitlab-uploads': gitlab_uploads,
+    'rm-watermark': rm_watermark
+  } = opts;
 
   const path = opts._[0];
   let buffer;
   if (data) buffer = Buffer.from(data, 'binary');
 
   const cml = new CML(opts);
-  const output = await cml.publish({ buffer, path, gitlab_uploads, ...opts });
+  const output = await cml.publish({
+    ...opts,
+    buffer,
+    path,
+    gitlab_uploads,
+    rm_watermark
+  });
 
   if (!file) print(output);
   else await fs.writeFile(file, output);
@@ -44,6 +55,8 @@ const argv = yargs
     'native',
     "Uses driver's native capabilities to upload assets instead of CML's storage."
   )
+  .boolean('rm-watermark')
+  .describe('rm-watermark', 'Avoid CML watermark.')
   .default('file')
   .describe(
     'file',
diff --git a/bin/cml-send-comment.js b/bin/cml-send-comment.js
index fecd657..60aaf08 100644
--- a/bin/cml-send-comment.js
+++ b/bin/cml-send-comment.js
@@ -36,7 +36,7 @@ const argv = yargs
   .deprecateOption('head-sha', 'Use commit-sha instead')
   .boolean('rm-watermark')
   .describe(
-    'no-watermark',
+    'rm-watermark',
     'Avoid watermark. CML needs a watermark to be able to distinguish CML reports from other comments in order to provide extra functionality.'
   )
   .default('repo')
diff --git a/bin/cml-tensorboard-dev.js b/bin/cml-tensorboard-dev.js
index b9134e3..c24b6b9 100644
--- a/bin/cml-tensorboard-dev.js
+++ b/bin/cml-tensorboard-dev.js
@@ -9,7 +9,7 @@ const { spawn } = require('child_process');
 const { homedir } = require('os');
 const tempy = require('tempy');
 
-const { exec } = require('../src/utils');
+const { exec, watermark_uri } = require('../src/utils');
 
 const { TB_CREDENTIALS } = process.env;
 
@@ -29,7 +29,8 @@ const run = async (opts) => {
     logdir,
     name,
     description,
-    title
+    title,
+    'rm-watermark': rm_watermark
   } = opts;
 
   // set credentials
@@ -70,6 +71,8 @@ const run = async (opts) => {
     if (matches.length) {
       let output = matches[0];
 
+      if (!rm_watermark) output = watermark_uri({ uri: output, type: 'tb' });
+
       if (md) output = `[${title || name}](${output})`;
 
       if (!file) print(output);
@@ -119,6 +122,7 @@ const argv = yargs
     'file',
     'Append the output to the given file. Create it if does not exist.'
   )
+  .describe('rm-watermark', 'Avoid CML watermark.')
   .alias('file', 'f')
   .help('h').argv;
 
diff --git a/package-lock.json b/package-lock.json
index 8b8a6fe..548474c 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,6 +1,6 @@
 {
   "name": "@dvcorg/cml",
-  "version": "0.2.1",
+  "version": "0.2.2",
   "lockfileVersion": 1,
   "requires": true,
   "dependencies": {
diff --git a/package.json b/package.json
index 806bc0e..cafc2e0 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@dvcorg/cml",
-  "version": "0.2.1",
+  "version": "0.2.2",
   "author": {
     "name": "DVC",
     "url": "http://cml.dev"
diff --git a/src/cml.js b/src/cml.js
index cfda4eb..a983047 100644
--- a/src/cml.js
+++ b/src/cml.js
@@ -4,7 +4,7 @@ const strip_auth = require('strip-url-auth');
 
 const Gitlab = require('./drivers/gitlab');
 const Github = require('./drivers/github');
-const { upload, exec } = require('./utils');
+const { upload, exec, watermark_uri } = require('./utils');
 
 const uri_no_trailing_slash = (uri) => {
   return uri.endsWith('/') ? uri.substr(0, uri.length - 1) : uri;
@@ -80,7 +80,7 @@ class CML {
   }
 
   async publish(opts = {}) {
-    const { title = '', md, native, gitlab_uploads } = opts;
+    const { title = '', md, native, gitlab_uploads, rm_watermark } = opts;
 
     let mime, uri;
     if (native || gitlab_uploads) {
@@ -90,6 +90,11 @@ class CML {
       ({ mime, uri } = await upload(opts));
     }
 
+    if (!rm_watermark) {
+      const [, type] = mime.split('/');
+      uri = watermark_uri({ uri, type });
+    }
+
     if (md && mime.match('(image|video)/.*'))
       return `![](${uri}${title ? ` "${title}"` : ''})`;
 
diff --git a/src/drivers/github.js b/src/drivers/github.js
index b96c08e..02270ef 100644
--- a/src/drivers/github.js
+++ b/src/drivers/github.js
@@ -1,3 +1,4 @@
+const url = require('url');
 const github = require('@actions/github');
 
 const CHECK_TITLE = 'CML Report';
@@ -17,10 +18,19 @@ const owner_repo = (opts) => {
   return { owner, repo };
 };
 
-const octokit = (token) => {
+const octokit = (token, repo) => {
   if (!token) throw new Error('token not found');
 
-  return github.getOctokit(token);
+  const octokit_options = {};
+
+  if (!repo.includes('github.com')) {
+    // GitHub Enterprise, use the: repo URL host + '/api/v3' - as baseURL
+    // as per: https://developer.github.com/enterprise/v3/enterprise-admin/#endpoint-urls
+    const { host } = new url.URL(repo);
+    octokit_options.baseUrl = `https://${host}/api/v3`;
+  }
+
+  return github.getOctokit(token, octokit_options);
 };
 
 class Github {
@@ -43,7 +53,8 @@ class Github {
     const { report: body, commit_sha } = opts;
 
     const { url: commit_url } = await octokit(
-      this.token
+      this.token,
+      this.repo
     ).repos.createCommitComment({
       ...owner_repo({ uri: this.repo }),
       body,
@@ -65,7 +76,7 @@ class Github {
     } = opts;
 
     const name = title;
-    return await octokit(this.token).checks.create({
+    return await octokit(this.token, this.repo).checks.create({
       ...owner_repo({ uri: this.repo }),
       head_sha,
       started_at,
@@ -83,7 +94,7 @@ class Github {
 
   async runner_token() {
     const { owner, repo } = owner_repo({ uri: this.repo });
-    const { actions } = octokit(this.token);
+    const { actions } = octokit(this.token, this.repo);
 
     if (typeof repo !== 'undefined') {
       const {
@@ -112,7 +123,7 @@ class Github {
   async unregister_runner(opts) {
     const { name } = opts;
     const { owner, repo } = owner_repo({ uri: this.repo });
-    const { actions } = octokit(this.token);
+    const { actions } = octokit(this.token, this.repo);
     const { id: runner_id } = await this.runner_by_name({ name });
 
     if (typeof repo !== 'undefined') {
@@ -132,7 +143,7 @@ class Github {
   async runner_by_name(opts = {}) {
     const { name } = opts;
     const { owner, repo } = owner_repo({ uri: this.repo });
-    const { actions } = octokit(this.token);
+    const { actions } = octokit(this.token, this.repo);
     let runners = [];
 
     if (typeof repo !== 'undefined') {
diff --git a/src/utils.js b/src/utils.js
index 0e340a6..9ec232d 100644
--- a/src/utils.js
+++ b/src/utils.js
@@ -135,6 +135,14 @@ const parse_param_newline = (param) => {
   return param.replace(/\\n/g, '\n');
 };
 
+const watermark_uri = (opts = {}) => {
+  const { uri, type } = opts;
+  const url = new URL(uri);
+  url.searchParams.append('cml', type);
+
+  return url.toString();
+};
+
 exports.exec = exec;
 exports.fetch_upload_data = fetch_upload_data;
 exports.upload = upload;
@@ -143,3 +151,4 @@ exports.sleep = sleep;
 exports.is_proc_running = is_proc_running;
 exports.ssh_public_from_private_rsa = ssh_public_from_private_rsa;
 exports.parse_param_newline = parse_param_newline;
+exports.watermark_uri = watermark_uri;


Test patch:
diff --git a/bin/cml-publish.test.js b/bin/cml-publish.test.js
index 358fa29..ea9bae3 100644
--- a/bin/cml-publish.test.js
+++ b/bin/cml-publish.test.js
@@ -20,6 +20,7 @@ describe('CML e2e', () => {
                                           [deprecated: Use --native instead] [boolean]
         --native          Uses driver's native capabilities to upload assets instead
                           of CML's storage.                                  [boolean]
+        --rm-watermark    Avoid CML watermark.                               [boolean]
         --file, -f        Append the output to the given file. Create it if does not
                           exist.
         --repo            Specifies the repo to be used. If not specified is extracted
diff --git a/bin/cml-send-comment.test.js b/bin/cml-send-comment.test.js
index 4c8513f..9159921 100644
--- a/bin/cml-send-comment.test.js
+++ b/bin/cml-send-comment.test.js
@@ -23,9 +23,9 @@ describe('Comment integration tests', () => {
         --commit-sha    Commit SHA linked to this comment. Defaults to HEAD.
         --head-sha      Commit SHA linked to this comment. Defaults to HEAD.
                                                   [deprecated: Use commit-sha instead]
-        --no-watermark  Avoid watermark. CML needs a watermark to be able to
+        --rm-watermark  Avoid watermark. CML needs a watermark to be able to
                         distinguish CML reports from other comments in order to
-                        provide extra functionality.
+                        provide extra functionality.                         [boolean]
         --repo          Specifies the repo to be used. If not specified is extracted
                         from the CI ENV.
         --token         Personal access token to be used. If not specified in
diff --git a/bin/cml-tensorboard-dev.test.js b/bin/cml-tensorboard-dev.test.js
index 395d9c4..1042554 100644
--- a/bin/cml-tensorboard-dev.test.js
+++ b/bin/cml-tensorboard-dev.test.js
@@ -37,6 +37,7 @@ describe('CML e2e', () => {
         --title, -t        Markdown title, if not specified, param name will be used.
         --file, -f         Append the output to the given file. Create it if does not
                            exist.
+        --rm-watermark     Avoid CML watermark.
         -h                 Show help                                         [boolean]
         --plugins"
     `);
@@ -57,5 +58,6 @@ describe('CML e2e', () => {
 
     expect(is_running).toBe(true);
     expect(output.startsWith(`[${title}](https://`)).toBe(true);
+    expect(output.includes('cml=tb')).toBe(true);
   });
 });
diff --git a/src/cml.test.js b/src/cml.test.js
index 06696bf..dc24de4 100644
--- a/src/cml.test.js
+++ b/src/cml.test.js
@@ -31,6 +31,7 @@ describe('Github tests', () => {
     const output = await new CML().publish({ path });
 
     expect(output.startsWith('https://')).toBe(true);
+    expect(output.includes('cml=png')).toBe(true);
   });
 
   test('Publish image with markdown', async () => {
@@ -41,6 +42,7 @@ describe('Github tests', () => {
 
     expect(output.startsWith('![](https://')).toBe(true);
     expect(output.endsWith(` "${title}")`)).toBe(true);
+    expect(output.includes('cml=png')).toBe(true);
   });
 
   test('Publish a non image file in markdown', async () => {
@@ -51,6 +53,8 @@ describe('Github tests', () => {
 
     expect(output.startsWith(`[${title}](https://`)).toBe(true);
     expect(output.endsWith(')')).toBe(true);
+    console.log(output);
+    expect(output.includes('cml=pdf')).toBe(true);
   });
 
   test('Comment should succeed with a valid sha', async () => {
@@ -108,6 +112,7 @@ describe('Gitlab tests', () => {
     });
 
     expect(output.startsWith('https://')).toBe(true);
+    expect(output.includes('cml=png')).toBe(true);
   });
 
   test('Publish image using gl with markdown', async () => {
@@ -123,6 +128,7 @@ describe('Gitlab tests', () => {
 
     expect(output.startsWith('![](https://')).toBe(true);
     expect(output.endsWith(` "${title}")`)).toBe(true);
+    expect(output.includes('cml=png')).toBe(true);
   });
 
   test('Publish a non image file using gl in markdown', async () => {
@@ -138,6 +144,7 @@ describe('Gitlab tests', () => {
 
     expect(output.startsWith(`[${title}](https://`)).toBe(true);
     expect(output.endsWith(')')).toBe(true);
+    expect(output.includes('cml=pdf')).toBe(true);
   });
 
   test('Publish a non image file using native', async () => {
@@ -153,6 +160,7 @@ describe('Gitlab tests', () => {
 
     expect(output.startsWith(`[${title}](https://`)).toBe(true);
     expect(output.endsWith(')')).toBe(true);
+    expect(output.includes('cml=pdf')).toBe(true);
   });
 
   test('Publish should fail with an invalid driver', async () => {