refactor usePrimaryEntityLocation

packages/datatrak-web/src/utils/index.ts

@@ -5,7 +5,6 @@
 
 export { errorToast, successToast } from './toast';
 export { useIsMobileMediaQuery as useIsMobile } from './useIsMobileMediaQuery';
-export { useFromLocation } from './useFromLocation';
 export * from './date';
 export * from './detectDevice';
 export { gaEvent } from './ga';
@@ -14,4 +13,4 @@ export {
   getTaskFilterSetting,
   removeTaskFilterSetting,
 } from './taskFilterSettings';
-export { usePrimaryEntityLocation } from './usePrimaryEntityLocation';
+export { useFromLocation, usePrimaryEntityLocation } from './useLocationState';

packages/datatrak-web/src/utils/useLocationState.ts

@@ -0,0 +1,25 @@
+/*
+ * Tupaia
+ *  Copyright (c) 2017 - 2024 Beyond Essential Systems Pty Ltd
+ */
+import { useLocation } from 'react-router-dom';
+
+function hasProperty<T extends string>(state: unknown, property: T): state is Record<T, string> {
+  if (state !== null && typeof state === 'object' && property in state) {
+    return typeof (state as Record<T, unknown>)[property] === 'string';
+  }
+  return false;
+}
+
+function useLocationState(property: 'primaryEntityCode' | 'from'): string | undefined {
+  const location = useLocation();
+  return hasProperty(location.state, property) ? location.state[property] : undefined;
+}
+
+export function usePrimaryEntityLocation() {
+  return useLocationState('primaryEntityCode');
+}
+
+export function useFromLocation() {
+  return useLocationState('from');
+}