Tree.tsx

import Collapse from '@material-ui/core/Collapse';
import FormControlLabel from '@material-ui/core/FormControlLabel';
import IconButton from '@material-ui/core/IconButton';
import { createStyles, withStyles, WithStyles } from '@material-ui/core/styles';
import DownArrow from '@material-ui/icons/KeyboardArrowDown';
import UpArrow from '@material-ui/icons/KeyboardArrowUp';
import produce from 'immer';
import React from 'react';
import LegoCheckbox from '../Checkbox/Checkbox';

const styles = () =>
  createStyles({
    option: {
      display: 'flex',
      justifyContent: 'space-between'
    }
  });

// Options Format
/*
  {
    [key]: {
      name: name,
      key: key,
      checked: false,
      children: [key1, key2],
      firstLevel: true,
    },
    [key1]: {
      name: name1,
      key: key,
      checked: false,
      children: [],
      firstLevel: false,
    },
    [key2]: {
      name: name2,
      key: key,
      checked: false,
      children: [name3],
      firstLevel: false,
    },
    [name3]: {
      name: name3,
      key: key,
      checked: false,
      children: [],
      firstLevel: true,
    }
  },
*/

class Tree extends React.Component<TreeProps, ITreeState> {
  constructor(props: TreeProps) {
    super(props);

    this.state = {
      opened: []
    };
  }

  onClose = (option: string) => {
    const opened = this.state.opened;
    opened.splice(this.state.opened.indexOf(option), 1);

    this.setState({
      opened
    });
  };

  onOpen = (option: string) => {
    this.setState({
      opened: this.state.opened.concat([option])
    });
  };

  isAllChildrenSelected = (parent: IOption, options: IOptionsMap) => {
    let allSelected = true;
    parent.children.forEach(child => {
      allSelected = allSelected && options[child].checked;
    });

    return allSelected;
  };

  handleChange = (parent: IOption) => (
    event: React.ChangeEvent<HTMLInputElement>
  ) => {
    const options = produce(this.props.options, opts => {
      opts[event.target.value].checked = event.target.checked;

      // If this option has children, select them all.
      const children = opts[event.target.value].children;
      if (children.length) {
        children.forEach(child => {
          opts[child].checked = event.target.checked;
        });
      }

      // If all the siblings are selected, select it's parent.
      if (parent) {
        opts[parent.key].checked = this.isAllChildrenSelected(parent, opts);
      }
    });
    this.props.onChange(options);
  };

  renderOption = (parent: IOption, option: IOption, level: number) => {
    const name = option.name;
    const hasChildren = option.children.length > 0;
    const expanded = this.state.opened.indexOf(name) !== -1;

    return (
      <div style={{ marginLeft: (level + 1) * 10 }} key={name}>
        <div className={this.props.classes.option}>
          <FormControlLabel
            control={
              <LegoCheckbox
                legoType="yellow"
                checked={option.checked}
                onChange={this.handleChange(parent)}
                value={option.key}
              />
            }
            label={name.slice(0, 1).toUpperCase() + name.slice(1).toLowerCase()}
          />
          {hasChildren && (
            <div>
              {expanded ? (
                <IconButton onClick={() => this.onClose(name)}>
                  <UpArrow />
                </IconButton>
              ) : (
                <IconButton onClick={() => this.onOpen(name)}>
                  <DownArrow />
                </IconButton>
              )}
            </div>
          )}
        </div>
        <Collapse in={expanded}>
          {hasChildren && (
            <div>
              {option.children.map(child =>
                this.renderOption(option, this.props.options[child], level + 1)
              )}
            </div>
          )}
        </Collapse>
      </div>
    );
  };

  render() {
    return (
      <div>
        {Object.keys(this.props.options)
          .map(k => this.props.options[k])
          .filter(option => option.firstLevel)
          .map(option => this.renderOption(undefined, option, 0))}
      </div>
    );
  }
}

interface ITreeProps {
  options: IOptionsMap;
  onChange: (options: any) => void;
}

interface IOptionsMap {
  [key: string]: IOption;
}

interface IOption {
  name: string;
  key: string;
  checked: boolean;
  children: string[];
  firstLevel: boolean;
}

interface ITreeState {
  opened: string[];
}

type TreeProps = ITreeProps & WithStyles<typeof styles>;

export default withStyles(styles)(Tree);