[theme] Correct augmentColor TypeScript definition

packages/material-ui/src/styles/createPalette.d.ts

@@ -1,6 +1,8 @@
 import { Color, PaletteType } from '..';
 import { CommonColors } from '../colors/common';
 
+export type ColorPartial = Partial<Color>;
+
 export interface TypeText {
   primary: string;
   secondary: string;
@@ -24,7 +26,7 @@ export interface TypeBackground {
 
 export type TypeDivider = string;
 
-export type PaletteColorOptions = SimplePaletteColorOptions | Partial<Color>;
+export type PaletteColorOptions = SimplePaletteColorOptions | ColorPartial;
 
 export interface SimplePaletteColorOptions {
   light?: string;
@@ -64,16 +66,18 @@ export interface Palette {
   action: TypeAction;
   background: TypeBackground;
   getContrastText: (background: string) => string;
-  augmentColor: (
-    color: SimplePaletteColorOptions,
-    mainShade: number | string,
-    lightShade: number | string,
-    darkShade: number | string,
-  ) => void;
+  augmentColor: {
+    (
+      color: ColorPartial,
+      mainShade?: number | string,
+      lightShade?: number | string,
+      darkShade?: number | string,
+    ): void;
+    (color: PaletteColorOptions): void;
+  };
 }
 
 export type PartialTypeObject = { [P in keyof TypeObject]?: Partial<TypeObject[P]> };
-export type ColorPartial = Partial<Color>;
 
 export interface PaletteOptions {
   primary?: PaletteColorOptions;

packages/material-ui/src/styles/createPalette.js

@@ -144,6 +144,8 @@ export default function createPalette(palette) {
     if (!color.contrastText) {
       color.contrastText = getContrastText(color.main);
     }
+
+    return color;
   }
 
   augmentColor(primary);

packages/material-ui/src/styles/createPalette.spec.ts

@@ -0,0 +1,22 @@
+import { Color } from '@material-ui/core';
+import blue from '@material-ui/core/colors/blue';
+import createPalette, {
+  PaletteColorOptions,
+  SimplePaletteColorOptions,
+} from '@material-ui/core/styles/createPalette';
+
+{
+  const palette = createPalette({});
+  const color: Color = blue;
+  const option: SimplePaletteColorOptions = { main: blue[400] };
+  const colorOrOption: PaletteColorOptions = undefined as any;
+
+  palette.augmentColor(color);
+  palette.augmentColor(color, 400);
+  palette.augmentColor(color, 400, 200, 600);
+  palette.augmentColor(color, 400, undefined, 600);
+  palette.augmentColor(option);
+  palette.augmentColor(option, 400); // $ExpectError
+  palette.augmentColor(colorOrOption);
+  palette.augmentColor(colorOrOption, 400); // $ExpectError
+}

packages/material-ui/src/styles/createPalette.test.js

@@ -361,11 +361,65 @@ describe('createPalette()', () => {
   });
 
   describe('augmentColor', () => {
+    const palette = createPalette({});
+
     it('should throw when the input is invalid', () => {
-      const palette = createPalette({});
       assert.throws(() => {
         palette.augmentColor({});
       }, /The color object needs to have a/);
     });
+
+    it('should accept a color', () => {
+      const color1 = palette.augmentColor({ ...indigo });
+      assert.deepEqual(
+        {
+          main: color1.main,
+          light: color1.light,
+          dark: color1.dark,
+          contrastText: color1.contrastText,
+        },
+        {
+          dark: '#303f9f',
+          light: '#7986cb',
+          main: '#3f51b5',
+          contrastText: '#fff',
+        },
+      );
+      const color2 = palette.augmentColor({ ...indigo }, 400, 200, 600);
+      assert.deepEqual(
+        {
+          light: color2.light,
+          main: color2.main,
+          dark: color2.dark,
+          contrastText: color2.contrastText,
+        },
+        {
+          light: '#9fa8da',
+          main: '#5c6bc0',
+          dark: '#3949ab',
+          contrastText: '#fff',
+        },
+      );
+    });
+
+    it('should accept a partial palette color', () => {
+      const color = palette.augmentColor({
+        main: indigo[500],
+      });
+      assert.deepEqual(
+        {
+          light: color.light,
+          main: color.main,
+          dark: color.dark,
+          contrastText: color.contrastText,
+        },
+        {
+          light: 'rgb(101, 115, 195)',
+          main: '#3f51b5',
+          dark: 'rgb(44, 56, 126)',
+          contrastText: '#fff',
+        },
+      );
+    });
   });
 });