Rewrite to React Hooks

wojtekmaj · Apr 29, 2024 · bb90c67 · bb90c67
1 parent 4888f1c
commit bb90c67
Showing 1 changed file with 70 additions and 62 deletions.
diff --git a/packages/react-fit/src/Fit.tsx b/packages/react-fit/src/Fit.tsx
@@ -1,6 +1,6 @@
 'use client';
 
-import { Children, Component } from 'react';
+import { Children, useCallback, useEffect, useRef } from 'react';
 import detectElementOverflow from 'detect-element-overflow';
 import warning from 'warning';
 
@@ -212,50 +212,40 @@ export type FitProps = {
   spacing?: number | Spacing;
 };
 
-export default class Fit extends Component<FitProps> {
-  componentDidMount() {
-    this.fit();
-
-    const onMutation = () => {
-      this.fit();
-    };
-
-    if (isMutationObserverSupported && this.element) {
-      const mutationObserver = new MutationObserver(onMutation);
-
-      mutationObserver.observe(this.element, {
-        attributes: true,
-        attributeFilter: ['class', 'style'],
-      });
-    }
-  }
-
-  container?: HTMLElement | null;
-  element?: HTMLElement | null;
-  elementWidth?: number;
-  elementHeight?: number;
-  scrollContainer?: HTMLElement;
-
-  fit = () => {
-    const { scrollContainer, container, element } = this;
-
-    if (!scrollContainer || !container || !element) {
+export default function Fit({
+  children,
+  invertAxis,
+  invertSecondaryAxis,
+  mainAxis = 'y',
+  spacing = 8,
+}: FitProps) {
+  const container = useRef<HTMLElement | undefined>(undefined);
+  const element = useRef<HTMLElement | undefined>(undefined);
+  const elementWidth = useRef<number | undefined>(undefined);
+  const elementHeight = useRef<number | undefined>(undefined);
+  const scrollContainer = useRef<HTMLElement | undefined>(undefined);
+
+  const fit = useCallback(() => {
+    if (!scrollContainer.current || !container.current || !element.current) {
       return;
     }
 
-    const elementWidth = element.clientWidth;
-    const elementHeight = element.clientHeight;
+    const currentElementWidth = element.current.clientWidth;
+    const currentElementHeight = element.current.clientHeight;
 
     // No need to recalculate - already did that for current dimensions
-    if (this.elementWidth === elementWidth && this.elementHeight === elementHeight) {
+    if (
+      elementWidth.current === currentElementWidth &&
+      elementHeight.current === currentElementHeight
+    ) {
       return;
     }
 
     // Save the dimensions so that we know we don't need to repeat the function if unchanged
-    this.elementWidth = elementWidth;
-    this.elementHeight = elementHeight;
+    elementWidth.current = currentElementWidth;
+    elementHeight.current = currentElementHeight;
 
-    const parent = container.parentElement;
+    const parent = container.current.parentElement;
 
     // Container was unmounted
     if (!parent) {
@@ -266,11 +256,11 @@ export default class Fit extends Component<FitProps> {
      * We need to ensure that <Fit />'s child has a absolute position. Otherwise,
      * we wouldn't be able to place the child in the correct position.
      */
-    const style = window.getComputedStyle(element);
+    const style = window.getComputedStyle(element.current);
     const { position } = style;
 
     if (position !== 'absolute') {
-      element.style.position = 'absolute';
+      element.current.style.position = 'absolute';
     }
 
     /**
@@ -284,42 +274,60 @@ export default class Fit extends Component<FitProps> {
       parent.style.position = 'relative';
     }
 
-    const { invertAxis, invertSecondaryAxis, mainAxis = 'y', spacing = 8 } = this.props;
-
     alignBothAxis({
       axis: mainAxis,
-      container,
-      element,
+      container: container.current,
+      element: element.current,
       invertAxis,
       invertSecondaryAxis,
-      scrollContainer,
+      scrollContainer: scrollContainer.current,
       spacing,
     });
-  };
+  }, [invertAxis, invertSecondaryAxis, mainAxis, spacing]);
+
+  const child = Children.only(children);
 
-  render() {
-    const { children } = this.props;
+  useEffect(() => {
+    fit();
 
-    const child = Children.only(children);
+    function onMutation() {
+      fit();
+    }
 
-    return (
-      <span
-        ref={(container) => {
-          this.container = container;
+    if (isMutationObserverSupported && element.current) {
+      const mutationObserver = new MutationObserver(onMutation);
 
-          const element = container && container.firstElementChild;
+      mutationObserver.observe(element.current, {
+        attributes: true,
+        attributeFilter: ['class', 'style'],
+      });
+    }
+  }, [fit]);
 
-          if (!element || !(element instanceof HTMLElement)) {
-            return;
-          }
+  function assignRefs(domElement: Element | null) {
+    if (!domElement || !(domElement instanceof HTMLElement)) {
+      return;
+    }
 
-          this.element = element;
-          this.scrollContainer = findScrollContainer(element);
-        }}
-        style={{ display: 'contents' }}
-      >
-        {child}
-      </span>
-    );
+    element.current = domElement;
+    scrollContainer.current = findScrollContainer(domElement);
   }
+
+  return (
+    <span
+      ref={(domContainer) => {
+        if (!domContainer) {
+          return;
+        }
+
+        container.current = domContainer;
+        const domElement = domContainer?.firstElementChild;
+
+        assignRefs(domElement);
+      }}
+      style={{ display: 'contents' }}
+    >
+      {child}
+    </span>
+  );
 }