Replace If Statements with Array.find

Altrim Beqiri

Altrim Beqiri /

Going through some of my code I found a simple function that returns a color based on a certain stock amount that looked similar to this

enum AlertColor {
  Red = "red",
  Yellow = "yellow",
  Green = "green",
}
const getColorForStockAmount = (stock = 0): AlertColor => {
  if (stock <= 100) {
    return AlertColor.Red;
  }
  if (stock <= 500) {
    return AlertColor.Yellow;
  }
  return AlertColor.Green;
};

getColorForStockAmount(50); // returns red
getColorForStockAmount(250); // returns yellow
getColorForStockAmount(750); // returns green

Now there is nothing wrong with the above snippet, it works pretty fine when you have few statements like this. It's just that when possible I like to avoid if statements, for example if you need to add additional checks it becomes a bit verbose for each condition we add. As you can see from the following diff we have 3 lines changed by adding another if statement. Imagine you had 5, 10 or more conditions you need to check, you can notice it could quickly get out of hand.

 if (stock <= 100) {
    return AlertColor.Red;
  }
+ if (stock <= 300) {           // Add line
+   return AlertColor.Orange;  // Add line
+ }                           // Add line
  if (stock <= 500) {
    return AlertColor.Yellow;
  }

Instead what you can do is have an array that holds the objects with value and color for each option and then use Array.find to get the item based on the condition met.

enum AlertColor {
  Red = "red",
  Yellow = "yellow",
  Green = "green",
}
type Amount = {
  amount: number;
  color: AlertColor;
};
const getColorForStockAmount = (stock = 0): AlertColor => {
  // NOTE: Since find returns the first element that satisfies the condition
  // it's important to have the objects ordered by the key you want to check
  const items: Amount[] = [
    { amount: 100, color: AlertColor.Red },
    { amount: 500, color: AlertColor.Yellow },
    { amount: Number.MAX_SAFE_INTEGER, color: AlertColor.Green },
  ];

  // .find() will return the first element that returns true for stock <= amount
  const item = items.find((item) => stock <= item.amount);

  return item?.color;
};

getColorForStockAmount(50); // returns red
getColorForStockAmount(250); // returns yellow
getColorForStockAmount(750); // returns green

What I like about this approach is that when we need another check we can easily add additional objects into the array without changing the surrounding code. As you can see in the following diff we have only one line changed and nothing else. The same would happen if we need to add 5, 10 or more objects, at maximum we would have one line change per item.

const items: Amount[] = [
  { amount: 100, color: AlertColor.Red },
  { amount: 300, color: AlertColor.Orange }, // Add line
  { amount: 500, color: AlertColor.Yellow },
  { amount: Number.MAX_SAFE_INTEGER, color: AlertColor.Green },
];

Moreover, you can easily use the same approach if for example you get the data from an API. Here we have a basic example usage in a React component where we fetch the items from the API and display the color in a badge.

enum AlertColor {
  Red = "red",
  Yellow = "yellow",
  Green = "green",
}
type Amount = {
  amount: number;
  color: AlertColor;
};
const StockBadge: React.FC<{ stock: number }> = ({ stock = 0 }) => {
  const [color, setColor] = useState<AlertColor>();

  useEffect(() => {
    // NOTE: Depending on the use case in a real world scenario we would
    // most likely do the fetching outside of this component and pass the items as a prop
    // so we avoid requesting the same API data multiple times
    const getColorForStockAmount = async (stock: number) => {
      // We fetch the items from our API
      const { data: items } = await http.get<Amount[]>("/stock");

      // We check the stock against the amount to get the proper item
      const item = items.find((item) => stock <= item.amount);

      // Finally we update the state with the item color
      setColor(item?.color);
    };

    getColorForStockAmount(stock);
  }, [stock]);

  return <div style={{ background: color }}>{stock}</div>;
};

<StockBadge stock={50} /> // renders <div style="background: red;">50</div>
<StockBadge stock={250} /> // renders <div style="background: yellow;">250</div>
<StockBadge stock={750} /> // renders <div style="background: green;">750</div>