All files / utils PluginEnvironmentProvider.tsx

93.54% Statements 29/31
66.66% Branches 6/9
100% Functions 10/10
92.85% Lines 26/28

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81  2x   2x 2x   2x 2x 2x 2x           30x       12x       10x       10x       5x   5x           5x 4x     1x         1x               2x               5x 5x   5x 5x       32x 32x     32x     2x  
import { JupyterFrontEnd } from '@jupyterlab/application';
import React, { createContext, useContext, useState } from 'react';
 
export const STUDIO_SAGEMAKER_UI_PLUGIN_ID = '@amzn/sagemaker-ui:project';
export const JUPYTERLAB_SAGEMAKER_UI_PLUGIN_ID = '@amzn/sagemaker-jupyterlab-extensions:sessionmanagement';
 
enum PluginEnvironmentType {
  LocalJL = 'local-jupyter-lab',
  JupyterLab = 'jupyterlab',
  Studio = 'studio',
}
 
class PluginEnvironment {
 
  public get isStudio(): boolean {
    return this.type === PluginEnvironmentType.Studio;
  }
 
  public get isLocalJL(): boolean {
    return this.type === PluginEnvironmentType.LocalJL;
  }
 
  public get isJupyterLab(): boolean {
    return this.type === PluginEnvironmentType.JupyterLab;
  }
 
  public get isStudioOrJupyterLab(): boolean {
    return this.isStudio || this.isJupyterLab;
  }
 
  constructor(
    public readonly type: PluginEnvironmentType,
  ) {
    console.debug(`PluginEnvironment created with type: ${type}`);
   }
 
}
 
function getPluginEnvironment(app: JupyterFrontEnd): PluginEnvironment {
  if (app.hasPlugin(STUDIO_SAGEMAKER_UI_PLUGIN_ID)) {
    return new PluginEnvironment(PluginEnvironmentType.Studio);
  }
 
  Iif (app.hasPlugin(JUPYTERLAB_SAGEMAKER_UI_PLUGIN_ID)
      || app.hasPlugin('@amzn/sagemaker-studio-scheduler:scheduler')) {
    return new PluginEnvironment(PluginEnvironmentType.JupyterLab);
  }
 
  return new PluginEnvironment(PluginEnvironmentType.LocalJL);
}
 
type PluginEnvironmentValue = {
  pluginEnvironment: PluginEnvironment;
  setPluginEnvironment: (state: PluginEnvironment) => void;
}
 
const PluginEnvironmentContext = createContext<PluginEnvironmentValue | undefined>(undefined);
 
type PluginEnvironmentProviderProps = {
  app: JupyterFrontEnd;
  children: React.ReactNode;
};
 
function PluginEnvironmentProvider({ app, children }: PluginEnvironmentProviderProps) {
  const [pluginEnvironment, setPluginEnvironment] = useState<PluginEnvironment>(() => {
    return getPluginEnvironment(app);
  });
  const value = { pluginEnvironment, setPluginEnvironment };
  return <PluginEnvironmentContext.Provider value={value}>{children}</PluginEnvironmentContext.Provider>
}
 
function usePluginEnvironment(): PluginEnvironmentValue {
  const context = useContext(PluginEnvironmentContext);
  Iif (context === undefined) {
    throw new Error('usePluginEnvironment must be used within a PluginEnvironmentProvider');
  }
  return context;
}
 
export { PluginEnvironmentProvider, usePluginEnvironment, PluginEnvironment, PluginEnvironmentType };